Fix the old flatten method which use the size(0) to caculate the batc…

…h size, the old method will intruduce Gather opertion in the onnx output, which will faild parsed by tensorRT 5.0 (#1134)
pytorch · Jul 19, 2019 · 2cae950 · 2cae950
1 parent bbd363c
commit 2cae950
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 9 deletions.
diff --git a/torchvision/models/alexnet.py b/torchvision/models/alexnet.py
@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 from .utils import load_state_dict_from_url
 
@@ -43,7 +44,7 @@ def __init__(self, num_classes=1000):
     def forward(self, x):
         x = self.features(x)
         x = self.avgpool(x)
-        x = x.view(x.size(0), 256 * 6 * 6)
+        x = torch.flatten(x, 1)
         x = self.classifier(x)
         return x
 

diff --git a/torchvision/models/densenet.py b/torchvision/models/densenet.py
@@ -154,7 +154,8 @@ def __init__(self, growth_rate=32, block_config=(6, 12, 24, 16),
     def forward(self, x):
         features = self.features(x)
         out = F.relu(features, inplace=True)
-        out = F.adaptive_avg_pool2d(out, (1, 1)).view(features.size(0), -1)
+        out = F.adaptive_avg_pool2d(out, (1, 1))
+        out = torch.flatten(out, 1)
         out = self.classifier(out)
         return out
 

diff --git a/torchvision/models/googlenet.py b/torchvision/models/googlenet.py
@@ -151,7 +151,7 @@ def forward(self, x):
 
         x = self.avgpool(x)
         # N x 1024 x 1 x 1
-        x = x.view(x.size(0), -1)
+        x = torch.flatten(x, 1)
         # N x 1024
         x = self.dropout(x)
         x = self.fc(x)
@@ -208,7 +208,7 @@ def forward(self, x):
         # aux1: N x 512 x 4 x 4, aux2: N x 528 x 4 x 4
         x = self.conv(x)
         # N x 128 x 4 x 4
-        x = x.view(x.size(0), -1)
+        x = torch.flatten(x, 1)
         # N x 2048
         x = F.relu(self.fc1(x), inplace=True)
         # N x 2048

diff --git a/torchvision/models/inception.py b/torchvision/models/inception.py
@@ -142,7 +142,7 @@ def forward(self, x):
         # N x 2048 x 1 x 1
         x = F.dropout(x, training=self.training)
         # N x 2048 x 1 x 1
-        x = x.view(x.size(0), -1)
+        x = torch.flatten(x, 1)
         # N x 2048
         x = self.fc(x)
         # N x 1000 (num_classes)
@@ -334,7 +334,7 @@ def forward(self, x):
         # Adaptive average pooling
         x = F.adaptive_avg_pool2d(x, (1, 1))
         # N x 768 x 1 x 1
-        x = x.view(x.size(0), -1)
+        x = torch.flatten(x, 1)
         # N x 768
         x = self.fc(x)
         # N x 1000

diff --git a/torchvision/models/resnet.py b/torchvision/models/resnet.py
@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 from .utils import load_state_dict_from_url
 
@@ -203,7 +204,7 @@ def forward(self, x):
         x = self.layer4(x)
 
         x = self.avgpool(x)
-        x = x.reshape(x.size(0), -1)
+        x = torch.flatten(x, 1)
         x = self.fc(x)
 
         return x

diff --git a/torchvision/models/squeezenet.py b/torchvision/models/squeezenet.py
@@ -99,7 +99,7 @@ def __init__(self, version='1_0', num_classes=1000):
     def forward(self, x):
         x = self.features(x)
         x = self.classifier(x)
-        return x.view(x.size(0), -1)
+        return torch.flatten(x, 1)
 
 
 def _squeezenet(version, pretrained, progress, **kwargs):

diff --git a/torchvision/models/vgg.py b/torchvision/models/vgg.py
@@ -1,3 +1,4 @@
+import torch
 import torch.nn as nn
 from .utils import load_state_dict_from_url
 
@@ -41,7 +42,7 @@ def __init__(self, features, num_classes=1000, init_weights=True):
     def forward(self, x):
         x = self.features(x)
         x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
+        x = torch.flatten(x, 1)
         x = self.classifier(x)
         return x