Merge branch 'dev/backbone_utils' into 'master'

Dev/backbone utils

See merge request open-mmlab/mmclassification!19

mmcls/models/utils/__init__.py

@@ -0,0 +1,4 @@
+from .channel_shuffle import channel_shuffle
+from .make_divisible import make_divisible
+
+__all__ = ['channel_shuffle', 'make_divisible']

mmcls/models/utils/channel_shuffle.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def channel_shuffle(x, groups):
+    """Channel Shuffle operation.
+
+    This function enables cross-group information flow for multiple groups
+    convolution layers.
+
+    Args:
+        x (Tensor): The input tensor.
+        groups (int): The number of groups to divide the input tensor
+            in the channel dimension.
+
+    Returns:
+        Tensor: The output tensor after channel shuffle operation.
+    """
+
+    batch_size, num_channels, height, width = x.size()
+    assert (num_channels % groups == 0), ('num_channels should be '
+                                          'divisible by groups')
+    channels_per_group = num_channels // groups
+
+    x = x.view(batch_size, groups, channels_per_group, height, width)
+    x = torch.transpose(x, 1, 2).contiguous()
+    x = x.view(batch_size, -1, height, width)
+
+    return x

mmcls/models/utils/make_divisible.py

@@ -0,0 +1,24 @@
+def make_divisible(value, divisor, min_value=None, min_ratio=0.9):
+    """Make divisible function.
+
+    This function rounds the channel number down to the nearest value that can
+    be divisible by the divisor.
+
+    Args:
+        value (int): The original channel number.
+        divisor (int): The divisor to fully divide the channel number.
+        min_value (int, optional): The minimum value of the output channel.
+            Default: None, means that the minimum value equal to the divisor.
+        min_ratio (float, optional): The minimum ratio of the rounded channel
+            number to the original channel number. Default: 0.9.
+    Returns:
+        int: The modified output channel number
+    """
+
+    if min_value is None:
+        min_value = divisor
+    new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than (1-min_ratio).
+    if new_value < min_ratio * value:
+        new_value += divisor
+    return new_value

tests/test_backbones/test_utils.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from mmcls.models.utils import channel_shuffle, make_divisible
+
+
+def test_make_divisible():
+    # test min_value is None
+    result = make_divisible(34, 8, None)
+    assert result == 32
+
+    # test when new_value > min_ratio * value
+    result = make_divisible(10, 8, min_ratio=0.9)
+    assert result == 16
+
+    # test min_value = 0.8
+    result = make_divisible(33, 8, min_ratio=0.8)
+    assert result == 32
+
+
+def test_channel_shuffle():
+    x = torch.randn(1, 24, 56, 56)
+    with pytest.raises(AssertionError):
+        # num_channels should be divisible by groups
+        channel_shuffle(x, 7)
+
+    groups = 3
+    batch_size, num_channels, height, width = x.size()
+    channels_per_group = num_channels // groups
+    out = channel_shuffle(x, groups)
+    # test the output value when groups = 3
+    for b in range(batch_size):
+        for c in range(num_channels):
+            c_out = c % channels_per_group * groups + c // channels_per_group
+            for i in range(height):
+                for j in range(width):
+                    assert x[b, c, i, j] == out[b, c_out, i, j]