Merge branch 'f/sandglass-block' into 'main'

Add sandglass block to embedded_vision_net

See merge request es/ai/hannah/hannah!375

hannah/callbacks/summaries.py

@@ -515,7 +515,6 @@ def __init__(self, module: torch.nn.Module):
     def run_node(self, n: torch.fx.Node):
         try:
             out = super().run_node(n)
-            print(out.shape, n)
         except Exception as e:
             print(str(e))
         if n.op == "call_function":
@@ -535,15 +534,15 @@ def run_node(self, n: torch.fx.Node):
                 self.data["MACs"] += [int(macs)]
             except Exception as e:
                 msglogger.warning("Summary of node %s failed: %s", n.name, str(e))
-                print(traceback.format_exc())   
+                print(traceback.format_exc())
         return out
 
 
 class FxMACSummaryCallback(MacSummaryCallback):
     def _do_summary(self, pl_module, input=None, print_log=True):
         interpreter = MACSummaryInterpreter(pl_module.model)
         dummy_input = input
-        
+
         if dummy_input is None:
             dummy_input = pl_module.example_feature_array
         dummy_input = dummy_input.to(pl_module.device)

hannah/models/embedded_vision_net/blocks.py

@@ -27,10 +27,10 @@ def grouped_pointwise(input, out_channels):
 
 @scope
 def expansion(input, expanded_channels):
-    #pw = partial(pointwise_conv2d, out_channels=expanded_channels)
-    #grouped_pw = partial(grouped_pointwise, out_channels=expanded_channels)
-    #return choice(input, pw, grouped_pw)
-    return pointwise_conv2d(input, out_channels=expanded_channels)
+    pw = partial(pointwise_conv2d, out_channels=expanded_channels)
+    grouped_pw = partial(grouped_pointwise, out_channels=expanded_channels)
+    return choice(input, pw, grouped_pw)
+    # return pointwise_conv2d(input, out_channels=expanded_channels)
 
 
 @scope
@@ -40,10 +40,10 @@ def spatial_correlation(input, out_channels, kernel_size, stride=1):
 
 @scope
 def reduction(input, out_channels):
-    #pw = partial(pointwise_conv2d, out_channels=out_channels)
-    #grouped_pw = partial(grouped_pointwise, out_channels=out_channels)
-    #return choice(input, pw, grouped_pw)
-    return pointwise_conv2d(input, out_channels=out_channels)
+    pw = partial(pointwise_conv2d, out_channels=out_channels)
+    grouped_pw = partial(grouped_pointwise, out_channels=out_channels)
+    return choice(input, pw, grouped_pw)
+    # return pointwise_conv2d(input, out_channels=out_channels)
 
 
 @scope
@@ -80,14 +80,31 @@ def expand_reduce(input, out_channels, expand_ratio, kernel_size, stride):
     return out
 
 
+# FIXME: integrate this into reduce_expand?
+@scope
+def sandglass_block(input, out_channels, reduce_ratio, kernel_size, stride):
+    in_channels = input.shape()[1]
+    reduced_channels = Int(in_channels / reduce_ratio)
+    out = depthwise_conv2d(input, out_channels=in_channels, kernel_size=kernel_size, stride=1)
+    out = batch_norm(out)
+    out = relu(out)
+    out = reduction(out, out_channels=reduced_channels)
+    out = expansion(out, expanded_channels=out_channels)
+    out = relu(out)
+    out = batch_norm(out)
+    out = depthwise_conv2d(out, out_channels=out_channels, kernel_size=kernel_size, stride=stride)
+    return out
+
+
 @scope
 def pattern(input, stride, out_channels, kernel_size, expand_ratio, reduce_ratio):
     convolution = partial(conv_relu, stride=stride, kernel_size=kernel_size, out_channels=out_channels)
     red_exp = partial(reduce_expand, out_channels=out_channels, reduce_ratio=reduce_ratio, kernel_size=kernel_size, stride=stride)
     exp_red = partial(expand_reduce, out_channels=out_channels, expand_ratio=expand_ratio, kernel_size=kernel_size, stride=stride)
-    #pool = partial(pooling, kernel_size=kernel_size, stride=stride)
+    pool = partial(pooling, kernel_size=kernel_size, stride=stride)
+    sandglass = partial(sandglass_block, out_channels=out_channels, reduce_ratio=reduce_ratio, kernel_size=kernel_size, stride=stride)
 
-    out = choice(input, convolution, exp_red, red_exp)
+    out = choice(input, convolution, exp_red, red_exp, pool, sandglass)
     return out
 
 

hannah/models/embedded_vision_net/operators.py

@@ -34,6 +34,8 @@ def grouped_conv2d(input, out_channels, kernel_size, stride, dilation=1, padding
     in_channels = input.shape()[1]
     if groups is None:
         groups = Groups(in_channels=in_channels, out_channels=out_channels, name="groups")
+    else:
+        groups = Int(groups)
 
     grouped_channels = Int(in_channels / groups)
     weight = Tensor(name='weight',

hannah/nas/search/search.py

@@ -71,7 +71,7 @@ def __init__(
             self.random_state = np.random.RandomState()
         else:
             self.random_state = random_state
-        
+
         self.example_input_array = None
         if input_shape is not None:
             self.example_input_array = torch.rand([1] + list(input_shape))
@@ -256,12 +256,9 @@ def build_model(self, parameters):
         return module
 
     def build_search_space(self):
-
-
         input = Tensor(
             "input", shape=self.example_input_array.shape, axis=("N", "C", "H", "W")
         )
-
         search_space = instantiate(self.config.model, input=input, _recursive_=True)
         return search_space