Merge pull request #122 from ekagra-ranjan/patch-1

Updated inceptionV4 to accept different sized images

Author: Cadene

pretrainedmodels/models/inceptionv4.py

@@ -1,6 +1,7 @@
 from __future__ import print_function, division, absolute_import
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torch.utils.model_zoo as model_zoo
 import os
 import sys
@@ -294,11 +295,12 @@ def __init__(self, num_classes=1001):
             Inception_C(),
             Inception_C()
         )
-        self.avg_pool = nn.AvgPool2d(8, count_include_pad=False)
         self.last_linear = nn.Linear(1536, num_classes)
 
     def logits(self, features):
-        x = self.avg_pool(features)
+        #Allows image of any size to be processed
+        adaptiveAvgPoolWidth = features.shape[2]
+        x = F.avg_pool2d(features, kernel_size=adaptiveAvgPoolWidth)
         x = x.view(x.size(0), -1)
         x = self.last_linear(x)
         return x
@@ -353,4 +355,4 @@ def inceptionv4(num_classes=1000, pretrained='imagenet'):
     print('success')
 
     # fail
-    assert inceptionv4(num_classes=1001, pretrained='imagenet')
+    assert inceptionv4(num_classes=1001, pretrained='imagenet')