Repository: pytorch/vision

Issue:
DenseNet memory_efficient mode appears subtly broken since torchscript changes
## 🐛 Bug

I was working with DenseNet recently, I noticed memory_efficient training was faster than with the flag disabled (unexpected), and it wasn't training very well at all, also errors from checkpoint fn about args not having required_grad=True set.

I believe #1342 subtly broke the checkpointing by wraping inputs in an extra layer of container

densenet.py:53
``` 
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(*inputs)

        return cp.checkpoint(closure, input)
```

should be

```
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(inputs)  # remove *

        return cp.checkpoint(closure, *input)  # add *
```
## To Reproduce

Train any densenet network with `memory_efficient=True`

## Expected behavior


## Environment

Any

 - PyTorch / torchvision Version (e.g., 1.0 / 0.4.0): 1.4
 - OS (e.g., Linux): any



DenseNet memory_efficient mode appears subtly broken since torchscript changes
## 🐛 Bug

I was working with DenseNet recently, I noticed memory_efficient training was faster than with the flag disabled (unexpected), and it wasn't training very well at all, also errors from checkpoint fn about args not having required_grad=True set.

I believe #1342 subtly broke the checkpointing by wraping inputs in an extra layer of container

densenet.py:53
``` 
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(*inputs)

        return cp.checkpoint(closure, input)
```

should be

```
    def call_checkpoint_bottleneck(self, input):
        # type: (List[Tensor]) -> Tensor
        def closure(*inputs):
            return self.bn_function(inputs)  # remove *

        return cp.checkpoint(closure, *input)  # add *
```
## To Reproduce

Train any densenet network with `memory_efficient=True`

## Expected behavior


## Environment

Any

 - PyTorch / torchvision Version (e.g., 1.0 / 0.4.0): 1.4
 - OS (e.g., Linux): any





Discussion:
Note, I didn't test my fix with torchscript so I may have broken that, my fix results in the inputs passed to bn_function being a tuple instead of the annotated list type
@rwightman can you send a PR fixing this? This should be fine with torchscript, because we skip torchscript for that function (it doesn't support `checkpoint`.

Also, would there be another test we could add that would catch this? IIRC we already have a test for the checkpoint behavior, but maybe it is missing some corner cases
@fmassa sure, I'll get a PR in for that, will take a look at the model in broken state again and see if any tests come to mind...
Note, I didn't test my fix with torchscript so I may have broken that, my fix results in the inputs passed to bn_function being a tuple instead of the annotated list type
@rwightman can you send a PR fixing this? This should be fine with torchscript, because we skip torchscript for that function (it doesn't support `checkpoint`.

Also, would there be another test we could add that would catch this? IIRC we already have a test for the checkpoint behavior, but maybe it is missing some corner cases
@fmassa sure, I'll get a PR in for that, will take a look at the model in broken state again and see if any tests come to mind...

Resolving patch:
diff --git a/torchvision/models/densenet.py b/torchvision/models/densenet.py
--- a/torchvision/models/densenet.py
+++ b/torchvision/models/densenet.py
@@ -53,9 +53,9 @@ def any_requires_grad(self, input):
     def call_checkpoint_bottleneck(self, input):
         # type: (List[Tensor]) -> Tensor
         def closure(*inputs):
-            return self.bn_function(*inputs)
+            return self.bn_function(inputs)
 
-        return cp.checkpoint(closure, input)
+        return cp.checkpoint(closure, *input)
 
     @torch.jit._overload_method  # noqa: F811
     def forward(self, input):


Test patch:
diff --git a/test/test_models.py b/test/test_models.py
--- a/test/test_models.py
+++ b/test/test_models.py
@@ -201,9 +201,11 @@ def test_memory_efficient_densenet(self):
         for name in ['densenet121', 'densenet169', 'densenet201', 'densenet161']:
             model1 = models.__dict__[name](num_classes=50, memory_efficient=True)
             params = model1.state_dict()
+            num_params = sum([x.numel() for x in model1.parameters()])
             model1.eval()
             out1 = model1(x)
             out1.sum().backward()
+            num_grad = sum([x.grad.numel() for x in model1.parameters() if x.grad is not None])
 
             model2 = models.__dict__[name](num_classes=50, memory_efficient=False)
             model2.load_state_dict(params)
@@ -212,6 +214,7 @@ def test_memory_efficient_densenet(self):
 
             max_diff = (out1 - out2).abs().max()
 
+            self.assertTrue(num_params == num_grad)
             self.assertTrue(max_diff < 1e-5)
 
     def test_resnet_dilation(self):