Repository: airctic/icevision

Issue:
Can't save a full model using torch.save (at least with faster-RCNN)
It is not possible to save a full model using default settings of `torch.save` (see stack trace below). This is because of the implementation of `remove_internal_model_transforms`, which uses inner functions in its implementation. The default pickle module does not support inner functions.

Workaround: use the `dill` module instead, which does support inner functions.

Suggested fix: It does not look as if the internal functions are necessary. If there were moved to standard functions, then the default pickle module should work.
`torch.save(model, 'mod.pth', pickle_module=pickle)` causes an error.

`torch.save(model, 'mod.pth', pickle_module=dill)` is a workaround.

**To Reproduce**

`torch.save(model, 'mod1-full.pth', pickle_module=pickle)`
results in:

```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-12-50f3761f4f3c> in <module>
----> 1 torch.save(model, 'mod1-full.pth', pickle_module=pickle)

~/anaconda3/envs/dlm/lib/python3.8/site-packages/torch/serialization.py in save(obj, f, pickle_module, pickle_protocol, _use_new_zipfile_serialization)
    370         if _use_new_zipfile_serialization:
    371             with _open_zipfile_writer(opened_file) as opened_zipfile:
--> 372                 _save(obj, opened_zipfile, pickle_module, pickle_protocol)
    373                 return
    374         _legacy_save(obj, opened_file, pickle_module, pickle_protocol)

~/anaconda3/envs/dlm/lib/python3.8/site-packages/torch/serialization.py in _save(obj, zip_file, pickle_module, pickle_protocol)
    474     pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)
    475     pickler.persistent_id = persistent_id
--> 476     pickler.dump(obj)
    477     data_value = data_buf.getvalue()
    478     zip_file.write_record('data.pkl', data_value, len(data_value))

AttributeError: Can't pickle local object 'remove_internal_model_transforms.<locals>.noop_normalize'
```

Relevant definition:
```
def remove_internal_model_transforms(model: GeneralizedRCNN):
    def noop_normalize(image: Tensor) -> Tensor:
        return image

    def noop_resize(
        image: Tensor, target: Optional[Dict[str, Tensor]]
    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
        return image, target

    model.transform.normalize = noop_normalize
    model.transform.resize = noop_resize
```




Discussion:

Resolving patch:
diff --git a/icevision/models/torchvision/utils.py b/icevision/models/torchvision/utils.py
--- a/icevision/models/torchvision/utils.py
+++ b/icevision/models/torchvision/utils.py
@@ -9,17 +9,19 @@
 from torchvision.models.detection.generalized_rcnn import GeneralizedRCNN
 
 
-def remove_internal_model_transforms(model: GeneralizedRCNN):
-    def noop_normalize(image: Tensor) -> Tensor:
-        return image
+def _noop_normalize(image: Tensor) -> Tensor:
+    return image
+
 
-    def noop_resize(
-        image: Tensor, target: Optional[Dict[str, Tensor]]
-    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
-        return image, target
+def _noop_resize(
+    image: Tensor, target: Optional[Dict[str, Tensor]]
+) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+    return image, target
 
-    model.transform.normalize = noop_normalize
-    model.transform.resize = noop_resize
+
+def remove_internal_model_transforms(model: GeneralizedRCNN):
+    model.transform.normalize = _noop_normalize
+    model.transform.resize = _noop_resize
 
 
 def patch_param_groups(


Test patch: