From 58f8d6518f154faa20df2d46d247692b141a1a0d Mon Sep 17 00:00:00 2001 From: Lara Haidar-Ahmad Date: Mon, 7 Jan 2019 08:35:16 -0800 Subject: [PATCH] support squeeznet onnx export --- torchvision/models/squeezenet.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/torchvision/models/squeezenet.py b/torchvision/models/squeezenet.py index a3e51e3b953..3801275c540 100644 --- a/torchvision/models/squeezenet.py +++ b/torchvision/models/squeezenet.py @@ -48,29 +48,29 @@ def __init__(self, version=1.0, num_classes=1000): self.features = nn.Sequential( nn.Conv2d(3, 96, kernel_size=7, stride=2), nn.ReLU(inplace=True), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(96, 16, 64, 64), Fire(128, 16, 64, 64), Fire(128, 32, 128, 128), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(256, 32, 128, 128), Fire(256, 48, 192, 192), Fire(384, 48, 192, 192), Fire(384, 64, 256, 256), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(512, 64, 256, 256), ) else: self.features = nn.Sequential( nn.Conv2d(3, 64, kernel_size=3, stride=2), nn.ReLU(inplace=True), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(64, 16, 64, 64), Fire(128, 16, 64, 64), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(128, 32, 128, 128), Fire(256, 32, 128, 128), - nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True), + nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=False), Fire(256, 48, 192, 192), Fire(384, 48, 192, 192), Fire(384, 64, 256, 256), @@ -79,7 +79,6 @@ def __init__(self, version=1.0, num_classes=1000): # Final convolution is initialized differently form the rest final_conv = nn.Conv2d(512, self.num_classes, kernel_size=1) self.classifier = nn.Sequential( - nn.Dropout(p=0.5), final_conv, nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1))