Repository: open-mmlab/mmdetection

Issue:
with_cp
Hi, 

I found that in the description of ResNeXt, there is an option of `with_cp` which is claimed to use checkpoint to save some memory. However, I couldn't find the the use of cp in ResNeXt.

Besides, ResNet have supporrt of with_cp for bottleneck structure only, setting with_cp with basiblock would raise an AssertionError, which is not stated in the description. 
with_cp
Hi, 

I found that in the description of ResNeXt, there is an option of `with_cp` which is claimed to use checkpoint to save some memory. However, I couldn't find the the use of cp in ResNeXt.

Besides, ResNet have supporrt of with_cp for bottleneck structure only, setting with_cp with basiblock would raise an AssertionError, which is not stated in the description. 


Discussion:
Hi @zyc-ai 
`with_cp` is used in `forward()` [here](https://github.com/open-mmlab/mmdetection/blob/a59f274ae6bb6c9d9e7d326837724fde7a882013/mmdet/models/backbones/resnet.py#L277). 
`Bottleneck` of `ResNeXt` inherits from `Bottleneck` of `ResNet`, so it doesn't need to implement it again. 
`BasicBlock` is used in `ResNet18`, which is relatively small in terms of memory consumption. `with_cp` is usually not necessary. 
Btw, you are very welcomed to create a PR to add `with_cp` option for `BasicBlock`. 
> Hi @zyc-ai
> `with_cp` is used in `forward()` [here](https://github.com/open-mmlab/mmdetection/blob/a59f274ae6bb6c9d9e7d326837724fde7a882013/mmdet/models/backbones/resnet.py#L277).
> `Bottleneck` of `ResNeXt` inherits from `Bottleneck` of `ResNet`, so it doesn't need to implement it again.
> `BasicBlock` is used in `ResNet18`, which is relatively small in terms of memory consumption. `with_cp` is usually not necessary.
> Btw, you are very welcomed to create a PR to add `with_cp` option for `BasicBlock`.

Ahhh, I forgot about the inheritance! I think it could be easier if we just modify description and mention small resnet does not support with_cp?
Hi @zyc-ai 
`with_cp` is used in `forward()` [here](https://github.com/open-mmlab/mmdetection/blob/a59f274ae6bb6c9d9e7d326837724fde7a882013/mmdet/models/backbones/resnet.py#L277). 
`Bottleneck` of `ResNeXt` inherits from `Bottleneck` of `ResNet`, so it doesn't need to implement it again. 
`BasicBlock` is used in `ResNet18`, which is relatively small in terms of memory consumption. `with_cp` is usually not necessary. 
Btw, you are very welcomed to create a PR to add `with_cp` option for `BasicBlock`. 
> Hi @zyc-ai
> `with_cp` is used in `forward()` [here](https://github.com/open-mmlab/mmdetection/blob/a59f274ae6bb6c9d9e7d326837724fde7a882013/mmdet/models/backbones/resnet.py#L277).
> `Bottleneck` of `ResNeXt` inherits from `Bottleneck` of `ResNet`, so it doesn't need to implement it again.
> `BasicBlock` is used in `ResNet18`, which is relatively small in terms of memory consumption. `with_cp` is usually not necessary.
> Btw, you are very welcomed to create a PR to add `with_cp` option for `BasicBlock`.

Ahhh, I forgot about the inheritance! I think it could be easier if we just modify description and mention small resnet does not support with_cp?

Resolving patch:
diff --git a/mmdet/models/backbones/resnet.py b/mmdet/models/backbones/resnet.py
--- a/mmdet/models/backbones/resnet.py
+++ b/mmdet/models/backbones/resnet.py
@@ -51,7 +51,7 @@ def __init__(self,
         self.downsample = downsample
         self.stride = stride
         self.dilation = dilation
-        assert not with_cp
+        self.with_cp = with_cp
 
     @property
     def norm1(self):
@@ -62,19 +62,29 @@ def norm2(self):
         return getattr(self, self.norm2_name)
 
     def forward(self, x):
-        identity = x
 
-        out = self.conv1(x)
-        out = self.norm1(out)
-        out = self.relu(out)
+        def _inner_forward(x):
+            identity = x
 
-        out = self.conv2(out)
-        out = self.norm2(out)
+            out = self.conv1(x)
+            out = self.norm1(out)
+            out = self.relu(out)
 
-        if self.downsample is not None:
-            identity = self.downsample(x)
+            out = self.conv2(out)
+            out = self.norm2(out)
+
+            if self.downsample is not None:
+                identity = self.downsample(x)
+
+            out += identity
+
+            return out
+
+        if self.with_cp and x.requires_grad:
+            out = cp.checkpoint(_inner_forward, x)
+        else:
+            out = _inner_forward(x)
 
-        out += identity
         out = self.relu(out)
 
         return out


Test patch:
diff --git a/tests/test_backbone.py b/tests/test_backbone.py
--- a/tests/test_backbone.py
+++ b/tests/test_backbone.py
@@ -49,10 +49,6 @@ def check_norm_state(modules, train_state):
 
 def test_resnet_basic_block():
 
-    with pytest.raises(AssertionError):
-        # Not implemented yet.
-        BasicBlock(64, 64, with_cp=True)
-
     with pytest.raises(AssertionError):
         # Not implemented yet.
         dcn = dict(type='DCN', deformable_groups=1, fallback_on_stride=False)
@@ -93,6 +89,13 @@ def test_resnet_basic_block():
     x_out = block(x)
     assert x_out.shape == torch.Size([1, 64, 56, 56])
 
+    # Test BasicBlock with checkpoint forward
+    block = BasicBlock(64, 64, with_cp=True)
+    assert block.with_cp
+    x = torch.randn(1, 64, 56, 56)
+    x_out = block(x)
+    assert x_out.shape == torch.Size([1, 64, 56, 56])
+
 
 def test_resnet_bottleneck():
 
@@ -285,10 +288,6 @@ def test_resnet_backbone():
         # In ResNet: 1 <= num_stages <= 4
         ResNet(50, num_stages=0)
 
-    with pytest.raises(AssertionError):
-        # with checkpoint is not implemented in BasicBlock of ResNet18
-        ResNet(18, with_cp=True)
-
     with pytest.raises(AssertionError):
         # len(stage_with_dcn) == num_stages
         dcn = dict(type='DCN', deformable_groups=1, fallback_on_stride=False)
@@ -379,6 +378,12 @@ def test_resnet_backbone():
     assert feat[2].shape == torch.Size([1, 256, 14, 14])
     assert feat[3].shape == torch.Size([1, 512, 7, 7])
 
+    # Test ResNet18 with checkpoint forward
+    model = ResNet(18, with_cp=True)
+    for m in model.modules():
+        if is_block(m):
+            assert m.with_cp
+
     # Test ResNet50 with BatchNorm forward
     model = ResNet(50)
     for m in model.modules():