Repository: pytorch/vision

Issue:
Use torch.jit.scripts to save the quantized inception_v3 error
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

```python
import torch
import torchvision.models.quantization as models

model = models.inception_v3(pretrained=True, progress=True, quantize=False)
model.to('cpu')
model.eval()
   
# Fuse Conv, bn and relu
model.fuse_model()
torch.jit.save(torch.jit.script(model), 'inception_v3_quantized_float.pth')
```

```shell
Traceback (most recent call last):
  File "resnt18.py", line 193, in <module>
    torch.jit.save(torch.jit.script(myModel), saved_model_dir + scripted_float_model_file)
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/__init__.py", line 1255, in script
    return torch.jit._recursive.recursive_script(obj)
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 534, in recursive_script
    return create_script_module(nn_module, infer_methods_to_compile(nn_module))
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 296, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, cpp_module, stubs)
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 340, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 259, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 574, in compile_unbound_method
    create_methods_from_stubs(concrete_type, (stub,))
  File "/usr/local/lib/python3.6/dist-packages/torch/jit/_recursive.py", line 259, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError:
Module 'QuantizableInception3' has no attribute 'AuxLogits' :
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/inception.py", line 161
        aux_defined = self.training and self.aux_logits
        if aux_defined:
            aux = self.AuxLogits(x)
                  ~~~~~~~~~~~~~~ <--- HERE
        else:
            aux = None
'QuantizableInception3._forward' is being compiled since it was called from 'QuantizableInception3.forward'
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/quantization/inception.py", line 196
        x = self._transform_input(x)
        x = self.quant(x)
        x, aux = self._forward(x)
        ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        x = self.dequant(x)
        aux_defined = self.training and self.aux_logits

```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->
Successfully saved

## Environment

Collecting environment information...
PyTorch version: 1.4.0
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.1 LTS
GCC version: (Ubuntu 6.5.0-2ubuntu1~18.04) 6.5.0 20181026
CMake version: version 3.14.7

Python version: 3.6
Versions of relevant libraries:
[pip3] numpy==1.18.0
[pip3] pytorch2sg==0.1.0
[pip3] torch==1.4.0
[pip3] torchvision==0.5.0
[conda] Could not collect

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi,

Thanks for the report!

This has been fixed in for GoogleNet in https://github.com/pytorch/vision/pull/1949, but not for Inception I believe.

cc @eellison  

Resolving patch:
diff --git a/torchvision/models/inception.py b/torchvision/models/inception.py
--- a/torchvision/models/inception.py
+++ b/torchvision/models/inception.py
@@ -55,7 +55,7 @@ def inception_v3(pretrained: bool = False, progress: bool = True, **kwargs: Any)
         model.load_state_dict(state_dict)
         if not original_aux_logits:
             model.aux_logits = False
-            del model.AuxLogits
+            model.AuxLogits = None
         return model
 
     return Inception3(**kwargs)
@@ -108,6 +108,7 @@ def __init__(
         self.Mixed_6c = inception_c(768, channels_7x7=160)
         self.Mixed_6d = inception_c(768, channels_7x7=160)
         self.Mixed_6e = inception_c(768, channels_7x7=192)
+        self.AuxLogits: Optional[nn.Module] = None
         if aux_logits:
             self.AuxLogits = inception_aux(768, num_classes)
         self.Mixed_7a = inception_d(768)
@@ -170,11 +171,10 @@ def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor]]:
         # N x 768 x 17 x 17
         x = self.Mixed_6e(x)
         # N x 768 x 17 x 17
-        aux_defined = self.training and self.aux_logits
-        if aux_defined:
-            aux = self.AuxLogits(x)
-        else:
-            aux = None
+        aux = torch.jit.annotate(Optional[Tensor], None)
+        if self.AuxLogits is not None:
+            if self.training:
+                aux = self.AuxLogits(x)
         # N x 768 x 17 x 17
         x = self.Mixed_7a(x)
         # N x 1280 x 8 x 8
diff --git a/torchvision/models/quantization/inception.py b/torchvision/models/quantization/inception.py
--- a/torchvision/models/quantization/inception.py
+++ b/torchvision/models/quantization/inception.py
@@ -67,7 +67,7 @@ def inception_v3(pretrained=False, progress=True, quantize=False, **kwargs):
         if quantize:
             if not original_aux_logits:
                 model.aux_logits = False
-                del model.AuxLogits
+                model.AuxLogits = None
             model_url = quant_model_urls['inception_v3_google' + '_' + backend]
         else:
             model_url = inception_module.model_urls['inception_v3_google']
@@ -80,7 +80,7 @@ def inception_v3(pretrained=False, progress=True, quantize=False, **kwargs):
         if not quantize:
             if not original_aux_logits:
                 model.aux_logits = False
-                del model.AuxLogits
+                model.AuxLogits = None
     return model
 
 


Test patch:
diff --git a/test/test_models.py b/test/test_models.py
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -339,6 +339,18 @@ def get_gn(num_channels):
         self.assertFalse(any(isinstance(x, nn.BatchNorm2d) for x in model.modules()))
         self.assertTrue(any(isinstance(x, nn.GroupNorm) for x in model.modules()))
 
+    def test_inceptionv3_eval(self):
+        # replacement for models.inception_v3(pretrained=True) that does not download weights
+        kwargs = {}
+        kwargs['transform_input'] = True
+        kwargs['aux_logits'] = True
+        kwargs['init_weights'] = False
+        model = models.Inception3(**kwargs)
+        model.aux_logits = False
+        model.AuxLogits = None
+        m = torch.jit.script(model.eval())
+        self.checkModule(m, "inception_v3", torch.rand(1, 3, 299, 299))
+
     def test_fasterrcnn_double(self):
         model = models.detection.fasterrcnn_resnet50_fpn(num_classes=50, pretrained_backbone=False)
         model.double()