File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff line change 11from __future__ import print_function , division , absolute_import
22import torch
33import torch .nn as nn
4+ import torch .nn .functional as F
45import torch .utils .model_zoo as model_zoo
56import os
67import sys
@@ -294,11 +295,12 @@ def __init__(self, num_classes=1001):
294295 Inception_C (),
295296 Inception_C ()
296297 )
297- self .avg_pool = nn .AvgPool2d (8 , count_include_pad = False )
298298 self .last_linear = nn .Linear (1536 , num_classes )
299299
300300 def logits (self , features ):
301- x = self .avg_pool (features )
301+ #Allows image of any size to be processed
302+ adaptiveAvgPoolWidth = features .shape [2 ]
303+ x = F .avg_pool2d (features , kernel_size = adaptiveAvgPoolWidth )
302304 x = x .view (x .size (0 ), - 1 )
303305 x = self .last_linear (x )
304306 return x
@@ -353,4 +355,4 @@ def inceptionv4(num_classes=1000, pretrained='imagenet'):
353355 print ('success' )
354356
355357 # fail
356- assert inceptionv4 (num_classes = 1001 , pretrained = 'imagenet' )
358+ assert inceptionv4 (num_classes = 1001 , pretrained = 'imagenet' )
You can’t perform that action at this time.
0 commit comments