refactor mutables.ops.conv when searching groups & add test

Author: marsggbo

hyperbox/mutables/ops/conv.py

@@ -59,7 +59,16 @@ def __init__(
         _stride = stride.max_value if isinstance(stride, ValueSpace) else stride
         _padding = padding.max_value if isinstance(padding, ValueSpace) else padding
         _dilation = dilation.min_value if isinstance(dilation, ValueSpace) else dilation
-        _groups = groups.min_value if isinstance(groups, ValueSpace) else groups
+
+        _groups = groups
+        if isinstance(groups, ValueSpace):
+            if isinstance(in_channels, ValueSpace):
+                if groups is not in_channels:
+                    print('groups must be the same as in_channels when in_channels is ValueSpace')
+                    groups = in_channels
+                _groups = groups.max_value
+            else:
+                _groups = groups.min_value
 
         self.format_args(_kernel_size, _stride, _padding, _dilation)
         conv_kwargs = kwargs
@@ -199,7 +208,7 @@ def forward_conv(self, x):
             filters = filters[:out_channels, ...]
         if self.search_kernel_size:
             filters = self.transform_kernel_size(filters)
-        if self.search_groups:
+        if self.search_groups and not self.search_in_channel:
             filters = self.get_filters_by_groups(filters, in_channels, groups).contiguous()
         if self.auto_padding:
             kernel_size = filters.shape[2:]
@@ -222,15 +231,6 @@ def get_filters_by_groups(self, filters, in_channels, groups):
             start = part_id * sub_in_channels
             filter_crops.append(sub_filter[:, start:start + sub_in_channels, :, :])
         filters = torch.cat(filter_crops, dim=0)
-
-        # indices = []
-        # for i in range(groups):
-        #     part_id = i % sub_ratio
-        #     start = part_id * sub_in_channels
-        #     indices.extend(list(range(start, start + sub_in_channels)))
-        # print(f"groups={groups}, in_channels={in_channels}, indices={indices}")
-        # filters = filters[:, indices, :, :]
-        print(f"groups={groups}, in_channels={in_channels}, filters.shape={filters.shape}")
         return filters
 
     def transform_kernel_size(self, filters):        

tests/mutables/test_op_conv.py

@@ -5,7 +5,62 @@
 
 from hyperbox.utils.calc_model_size import flops_size_counter
 
-if __name__ == '__main__':
+def test_groups(cin, cout, group_list):
+    print(f"test_groups: cin={cin}, cout={cout}, group_list={group_list}\n")
+    groups = ValueSpace(candidates=group_list)
+    x = torch.rand(2, cin, 64, 64)
+    conv = ops.Conv2d(cin, cout, 3, 1, 1, groups=groups)
+    m = RandomMutator(conv)
+    for i in range(10):
+        m.reset()
+        print(f'\n*******step{i}********\n', conv)
+        y = conv(x)
+
+def test_cin_groups(cin_list, cout, group_list):
+    print(f"test_cin_groups: cin_list={cin_list}, cout={cout}, group_list={group_list}\n")
+    cin = ValueSpace(candidates=cin_list)
+    groups = ValueSpace(candidates=group_list)
+    x = torch.rand(2, 3, 64, 64)
+    conv = torch.nn.Sequential(
+        ops.Conv2d(3, cin, 3, 1, 1),
+        ops.Conv2d(cin, cout, 3, 1, 1, groups=cin)
+    )
+    m = RandomMutator(conv)
+    for i in range(10):
+        m.reset()
+        print(f'\n*******step{i}********\n', conv)
+        y = conv(x)
+
+def test_cout_groups(cin, cout_list, group_list):
+    print(f"test_cout_groups: cin={cin}, cout_list={cout_list}, group_list={group_list}\n")
+    cout = ValueSpace(candidates=cout_list)
+    groups = ValueSpace(candidates=group_list)
+    x = torch.rand(2, cin, 64, 64)
+    conv = ops.Conv2d(cin, cout, 3, 1, 1, groups=groups)
+    m = RandomMutator(conv)
+    for i in range(10):
+        m.reset()
+        print(f'\n*******step{i}********\n', conv)
+        y = conv(x)
+
+def test_cin_cout_groups(cin_list, cout_list, group_list):
+    print(f"test_cin_cout_groups: cin_list={cin_list}, cout_list={cout_list}, group_list={group_list}\n")
+    cin = ValueSpace(candidates=cin_list)
+    cout = ValueSpace(candidates=cout_list)
+    groups = ValueSpace(candidates=group_list)
+    x = torch.rand(2, 3, 64, 64)
+    conv = torch.nn.Sequential(
+        ops.Conv2d(3, cin, 3, 1, 1),
+        ops.Conv2d(cin, cout, 3, 1, 1, groups=cin)
+    )
+    m = RandomMutator(conv)
+    for i in range(10):
+        m.reset()
+        print(f'\n*******step{i}********\n', conv)
+        y = conv(x)
+
+def test_conv():
+    print('testing conv flops and sizes')
     x = torch.rand(1,3,64,64)
     vs1 = ValueSpace([10,2])
     vs2 = ValueSpace([3,5,7])
@@ -17,15 +72,22 @@
     )
     m = RandomMutator(op)
     m.reset()
-    print(op)
-    print(m._cache)
-    print(conv.weight.shape)
-    print(conv(x).shape)
-    r = flops_size_counter(op, (1,3,8,8), True, True)
-    print(r)
+    # print(op)
+    # print(m._cache)
+    # print(conv.weight.shape)
+    # print(conv(x).shape)
+    r = flops_size_counter(op, (1,3,8,8), False, False)
+    # print(r)
     op = torch.nn.Sequential(
         ops.Conv2d(3,8,3,1),
         ops.BatchNorm2d(8)
     )
-    r = flops_size_counter(op, (1,3,8,8), False, True)
-    print(conv(x).shape)
+    r = flops_size_counter(op, (1,3,8,8), False, False)
+    # print(conv(x).shape)
+
+if __name__ == '__main__':
+    test_conv()
+    test_groups(32, 64, group_list=[1, 2, 4, 8, 16, 32])
+    test_cin_groups([32, 64], 128, group_list=[32, 64])
+    test_cout_groups(8, [16, 32], group_list=[1, 2, 4, 8])
+    test_cin_cout_groups([8, 16], [32, 64], group_list=[1, 8, 16])