[NOMERGE] Dropblock ResNet50 training

docs/source/ops.rst

@@ -20,6 +20,8 @@ Operators
     box_iou
     clip_boxes_to_image
     deform_conv2d
+    drop_block2d
+    drop_block3d
     generalized_box_iou
     generalized_box_iou_loss
     masks_to_boxes
@@ -47,3 +49,5 @@ Operators
     FrozenBatchNorm2d
     ConvNormActivation
     SqueezeExcitation
+    DropBlock2d
+    DropBlock3d

references/classification/README.md

@@ -53,6 +53,13 @@ torchrun --nproc_per_node=8 train.py --model $MODEL
 
 Here `$MODEL` is one of `resnet18`, `resnet34`, `resnet50`, `resnet101` or `resnet152`.
 
+### ResNet with dropblock
+```
+torchrun --nproc_per_node=8 train.py --model resnet50 -b 128 --lr 0.4 --epochs 270
+```
+
+
+
 ### ResNext
 ```
 torchrun --nproc_per_node=8 train.py\

references/classification/train.py

@@ -288,6 +288,7 @@ def main(args):
             f"Invalid lr scheduler '{args.lr_scheduler}'. Only StepLR, CosineAnnealingLR and ExponentialLR "
             "are supported."
         )
+    main_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[125, 200, 250], gamma=0.1)
 
     if args.lr_warmup_epochs > 0:
         if args.lr_warmup_method == "linear":

test/test_ops.py

@@ -2,6 +2,7 @@
 import os
 from abc import ABC, abstractmethod
 from functools import lru_cache
+from itertools import product
 from typing import Callable, List, Tuple
 
 import numpy as np
@@ -57,6 +58,16 @@ def forward(self, a):
         self.layer(a)
 
 
+class DropBlockWrapper(nn.Module):
+    def __init__(self, obj):
+        super().__init__()
+        self.layer = obj
+        self.n_inputs = 1
+
+    def forward(self, a):
+        self.layer(a)
+
+
 class RoIOpTester(ABC):
     dtype = torch.float64
 
@@ -1357,5 +1368,87 @@ def test_split_normalization_params(self, norm_layer):
         assert len(params[1]) == 82
 
 
+class TestDropBlock:
+    @pytest.mark.parametrize("seed", range(10))
+    @pytest.mark.parametrize("dim", [2, 3])
+    @pytest.mark.parametrize("p", [0, 0.5])
+    @pytest.mark.parametrize("block_size", [5, 11])
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_drop_block(self, seed, dim, p, block_size, inplace):
+        torch.manual_seed(seed)
+        batch_size = 5
+        channels = 3
+        height = 11
+        width = height
+        depth = height
+        if dim == 2:
+            x = torch.ones(size=(batch_size, channels, height, width))
+            layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace)
+            feature_size = height * width
+        elif dim == 3:
+            x = torch.ones(size=(batch_size, channels, depth, height, width))
+            layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace)
+            feature_size = depth * height * width
+        layer.__repr__()
+
+        out = layer(x)
+        if p == 0:
+            assert out.equal(x)
+        if block_size == height:
+            for b, c in product(range(batch_size), range(channels)):
+                assert out[b, c].count_nonzero() in (0, feature_size)
+
+    @pytest.mark.parametrize("seed", range(10))
+    @pytest.mark.parametrize("dim", [2, 3])
+    @pytest.mark.parametrize("p", [0.1, 0.2])
+    @pytest.mark.parametrize("block_size", [3])
+    @pytest.mark.parametrize("inplace", [False])
+    def test_drop_block_random(self, seed, dim, p, block_size, inplace):
+        torch.manual_seed(seed)
+        batch_size = 5
+        channels = 3
+        height = 11
+        width = height
+        depth = height
+        if dim == 2:
+            x = torch.ones(size=(batch_size, channels, height, width))
+            layer = ops.DropBlock2d(p=p, block_size=block_size, inplace=inplace)
+        elif dim == 3:
+            x = torch.ones(size=(batch_size, channels, depth, height, width))
+            layer = ops.DropBlock3d(p=p, block_size=block_size, inplace=inplace)
+
+        trials = 250
+        num_samples = 0
+        counts = 0
+        cell_numel = torch.tensor(x.shape).prod()
+        for _ in range(trials):
+            with torch.no_grad():
+                out = layer(x)
+            non_zero_count = out.nonzero().size(0)
+            counts += cell_numel - non_zero_count
+            num_samples += cell_numel
+
+        assert abs(p - counts / num_samples) / p < 0.15
+
+    def make_obj(self, dim, p, block_size, inplace, wrap=False):
+        if dim == 2:
+            obj = ops.DropBlock2d(p, block_size, inplace)
+        elif dim == 3:
+            obj = ops.DropBlock3d(p, block_size, inplace)
+        return DropBlockWrapper(obj) if wrap else obj
+
+    @pytest.mark.parametrize("dim", (2, 3))
+    @pytest.mark.parametrize("p", [0, 1])
+    @pytest.mark.parametrize("block_size", [5, 7])
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_is_leaf_node(self, dim, p, block_size, inplace):
+        op_obj = self.make_obj(dim, p, block_size, inplace, wrap=True)
+        graph_node_names = get_graph_node_names(op_obj)
+
+        assert len(graph_node_names) == 2
+        assert len(graph_node_names[0]) == len(graph_node_names[1])
+        assert len(graph_node_names[0]) == 1 + op_obj.n_inputs
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

torchvision/models/resnet.py

@@ -6,6 +6,7 @@
 
 from .._internally_replaced_utils import load_state_dict_from_url
 from ..utils import _log_api_usage_once
+from ..ops import DropBlock2d
 
 
 __all__ = [
@@ -122,6 +123,7 @@ def __init__(
         base_width: int = 64,
         dilation: int = 1,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
+        p: float = 0.0,
     ) -> None:
         super().__init__()
         if norm_layer is None:
@@ -130,31 +132,40 @@ def __init__(
         # Both self.conv2 and self.downsample layers downsample the input when stride != 1
         self.conv1 = conv1x1(inplanes, width)
         self.bn1 = norm_layer(width)
+        # we won't be doing scheduled p
+        self.drop1 = DropBlock2d(p, 7)
         self.conv2 = conv3x3(width, width, stride, groups, dilation)
         self.bn2 = norm_layer(width)
+        self.drop2 = DropBlock2d(p, 7)
         self.conv3 = conv1x1(width, planes * self.expansion)
         self.bn3 = norm_layer(planes * self.expansion)
+        self.drop3 = DropBlock2d(p, 7)
         self.relu = nn.ReLU(inplace=True)
         self.downsample = downsample
+        self.drop4 = DropBlock2d(p, 7)
         self.stride = stride
 
     def forward(self, x: Tensor) -> Tensor:
         identity = x
-
+        # as in https://github.com/tensorflow/tpu/blob/b24729de804fdb751b06467d3dce0637fa652060/models/official/resnet/resnet_model.py#L545-L579
         out = self.conv1(x)
         out = self.bn1(out)
         out = self.relu(out)
+        out = self.drop1(out)
 
         out = self.conv2(out)
         out = self.bn2(out)
         out = self.relu(out)
+        out = self.drop2(out)
 
         out = self.conv3(out)
         out = self.bn3(out)
+        out = self.drop3(out)
 
         if self.downsample is not None:
             identity = self.downsample(x)
 
+        identity = self.drop4(identity)
         out += identity
         out = self.relu(out)
 
@@ -198,8 +209,9 @@ def __init__(
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
         self.layer1 = self._make_layer(block, 64, layers[0])
         self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
+        # https://github.com/tensorflow/tpu/blob/b24729de804fdb751b06467d3dce0637fa652060/models/official/resnet/resnet_main.py#L393-L394
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1], p=0.1 / 4)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2], p=0.1)
         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
         self.fc = nn.Linear(512 * block.expansion, num_classes)
 
@@ -227,6 +239,7 @@ def _make_layer(
         blocks: int,
         stride: int = 1,
         dilate: bool = False,
+        p: float = 0.0,
     ) -> nn.Sequential:
         norm_layer = self._norm_layer
         downsample = None
@@ -243,7 +256,7 @@ def _make_layer(
         layers = []
         layers.append(
             block(
-                self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer
+                self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer, p
             )
         )
         self.inplanes = planes * block.expansion
@@ -256,6 +269,7 @@ def _make_layer(
                     base_width=self.base_width,
                     dilation=self.dilation,
                     norm_layer=norm_layer,
+                    p=p
                 )
             )
 

torchvision/ops/__init__.py

@@ -11,6 +11,7 @@
 )
 from .boxes import box_convert
 from .deform_conv import deform_conv2d, DeformConv2d
+from .drop_block import drop_block2d, DropBlock2d, drop_block3d, DropBlock3d
 from .feature_pyramid_network import FeaturePyramidNetwork
 from .focal_loss import sigmoid_focal_loss
 from .giou_loss import generalized_box_iou_loss
@@ -54,4 +55,8 @@
     "ConvNormActivation",
     "SqueezeExcitation",
     "generalized_box_iou_loss",
+    "drop_block2d",
+    "DropBlock2d",
+    "drop_block3d",
+    "DropBlock3d",
 ]