add kws with transformers, and bitnet train script update

Author: Jzjerry

.gitignore

@@ -3,5 +3,5 @@ temp
 data
 
 .vscode
-
+*.log
 output/

MiCoCodeGen.py

@@ -497,8 +497,14 @@ def handle_call_method(self, n: torch.fx.node.Node, out: torch.Tensor):
             self.add_forward_call("MiCo_CONNECT", out, n.name, [src_name])
         elif method == "mean":
             src_name = input_names[0]
-            dim = self._resolve_arg_value(n.args[1])
-            keepdim = self._resolve_arg_value(n.args[2]) if len(n.args) > 2 else False
+            dim = self._resolve_arg_value(n.args[1]) if len(n.args) > 1 else self._resolve_arg_value(n.kwargs.get("dim", None))
+            keepdim = (
+                self._resolve_arg_value(n.args[2])
+                if len(n.args) > 2
+                else self._resolve_arg_value(n.kwargs.get("keepdim", False))
+            )
+            if dim is None:
+                raise NotImplementedError("Mean over all elements is not supported")
             self.add_uninitialized_tensor(n.name, out)
             if keepdim:
                 self.add_forward_call(f"MiCo_meankp{out.dim()}d_{{dtype}}", out, n.name, [src_name], [dim])

MiCoRegistry.py

@@ -233,24 +233,51 @@ def handle_tanh(codegen, n, out, input_names, input_args):
     codegen.add_forward_call("MiCo_tanh{dim}d_{dtype}", out, n.name, input_names)
 
 
+def _extract_scalar_param(param, param_name, default=None):
+    """Extract a scalar C API parameter from PyTorch int/tuple pooling args."""
+    if param is None:
+        if default is None:
+            raise ValueError(f"{param_name} cannot be None")
+        param = default
+
+    if isinstance(param, torch.fx.node.Node):
+        raise ValueError(f"Unresolved FX node for {param_name}: {param}")
+
+    if isinstance(param, torch.Size):
+        param = tuple(param)
+
+    if isinstance(param, (tuple, list)):
+        if len(param) == 0:
+            raise ValueError(f"{param_name} cannot be empty")
+        first = param[0]
+        if any(value != first for value in param):
+            raise NotImplementedError(
+                f"MiCo C pooling kernels only support scalar/symmetric {param_name}, got {param}"
+            )
+        param = first
+
+    if isinstance(param, bool) or not isinstance(param, int):
+        raise ValueError(f"Unexpected {param_name} type: {type(param)}")
+    return param
+
+
 def _extract_kernel_size(param):
-    """Helper to extract kernel size from tuple or int parameter."""
-    if isinstance(param, Tuple):
-        return param[0]
-    elif isinstance(param, int):
-        return param
-    else:
-        raise ValueError(f"Unexpected kernel_size type: {type(param)}")
+    """Helper to extract scalar kernel size for the C pooling API."""
+    return _extract_scalar_param(param, "kernel_size")
 
 
 def _extract_output_size(param):
-    """Helper to extract output size from tuple or int parameter."""
-    if isinstance(param, Tuple):
-        return param[0]
-    elif isinstance(param, int):
-        return param
+    """Helper to extract scalar output size for the C adaptive pooling API."""
+    return _extract_scalar_param(param, "output_size")
+
+
+def _pool_arg(n, input_args, index, name, default=None):
+    """Read pooling arg from positional or keyword FX args and normalize it."""
+    if len(input_args) > index:
+        value = input_args[index]
     else:
-        raise ValueError(f"Unexpected output_size type: {type(param)}")
+        value = n.kwargs.get(name, default)
+    return _extract_scalar_param(value, name, default)
 
 
 @MiCoOpRegistry.register_function(torch.nn.functional.linear)
@@ -275,20 +302,22 @@ def handle_linear(codegen, n, out, input_names, input_args):
 def handle_avg_pool2d(codegen, n, out, input_names, input_args):
     """Handler for 2D average pooling function."""
     codegen.add_uninitialized_tensor(n.name, out)
-    kernel_size = _extract_kernel_size(input_args[1])
-    stride = input_args[2] if len(input_args) > 2 else 1
+    kernel_size = _pool_arg(n, input_args, 1, "kernel_size")
+    stride = _pool_arg(n, input_args, 2, "stride", kernel_size)
+    padding = _pool_arg(n, input_args, 3, "padding", 0)
     codegen.add_forward_call("MiCo_avgpool{dim}d_{dtype}", out, n.name, input_names, 
-                             [kernel_size, stride])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_function(torch.nn.functional.max_pool2d)
 def handle_max_pool2d(codegen, n, out, input_names, input_args):
     """Handler for 2D max pooling function."""
     codegen.add_uninitialized_tensor(n.name, out)
-    kernel_size = _extract_kernel_size(input_args[1])
-    stride = input_args[2] if len(input_args) > 2 else 1
+    kernel_size = _pool_arg(n, input_args, 1, "kernel_size")
+    stride = _pool_arg(n, input_args, 2, "stride", kernel_size)
+    padding = _pool_arg(n, input_args, 3, "padding", 0)
     codegen.add_forward_call("MiCo_maxpool{dim}d_{dtype}", out, n.name, input_names,
-                             [kernel_size, stride])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_function(torch.nn.functional.adaptive_avg_pool2d)
@@ -303,20 +332,22 @@ def handle_adaptive_avg_pool2d(codegen, n, out, input_names, input_args):
 def handle_avg_pool1d(codegen, n, out, input_names, input_args):
     """Handler for 1D average pooling function."""
     codegen.add_uninitialized_tensor(n.name, out)
-    kernel_size = _extract_kernel_size(input_args[1])
-    stride = input_args[2] if len(input_args) > 2 else 1
+    kernel_size = _pool_arg(n, input_args, 1, "kernel_size")
+    stride = _pool_arg(n, input_args, 2, "stride", kernel_size)
+    padding = _pool_arg(n, input_args, 3, "padding", 0)
     codegen.add_forward_call("MiCo_avgpool{dim}d_{dtype}", out, n.name, input_names, 
-                             [kernel_size, stride])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_function(torch.nn.functional.max_pool1d)
 def handle_max_pool1d(codegen, n, out, input_names, input_args):
     """Handler for 1D max pooling function."""
     codegen.add_uninitialized_tensor(n.name, out)
-    kernel_size = _extract_kernel_size(input_args[1])
-    stride = input_args[2] if len(input_args) > 2 else 1
+    kernel_size = _pool_arg(n, input_args, 1, "kernel_size")
+    stride = _pool_arg(n, input_args, 2, "stride", kernel_size)
+    padding = _pool_arg(n, input_args, 3, "padding", 0)
     codegen.add_forward_call("MiCo_maxpool{dim}d_{dtype}", out, n.name, input_names,
-                             [kernel_size, stride])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_function(torch.nn.functional.adaptive_avg_pool1d)
@@ -550,8 +581,10 @@ def handle_avgpool2d_module(codegen, n, out, module, input_names):
     layer_name = n.name
     codegen.add_uninitialized_tensor(layer_name, out)
     kernel_size = _extract_kernel_size(module.kernel_size)
+    stride = _extract_scalar_param(module.stride, "stride", kernel_size)
+    padding = _extract_scalar_param(module.padding, "padding", 0)
     codegen.add_forward_call("MiCo_avgpool{dim}d_{dtype}", out, layer_name, input_names, 
-                             [kernel_size, module.stride, module.padding])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_module(torch.nn.MaxPool2d)
@@ -560,8 +593,10 @@ def handle_maxpool2d_module(codegen, n, out, module, input_names):
     layer_name = n.name
     codegen.add_uninitialized_tensor(layer_name, out)
     kernel_size = _extract_kernel_size(module.kernel_size)
+    stride = _extract_scalar_param(module.stride, "stride", kernel_size)
+    padding = _extract_scalar_param(module.padding, "padding", 0)
     codegen.add_forward_call("MiCo_maxpool{dim}d_{dtype}", out, layer_name, input_names, 
-                             [kernel_size, module.stride, module.padding])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_module(torch.nn.AdaptiveAvgPool2d)
@@ -579,8 +614,10 @@ def handle_avgpool1d_module(codegen, n, out, module, input_names):
     layer_name = n.name
     codegen.add_uninitialized_tensor(layer_name, out)
     kernel_size = _extract_kernel_size(module.kernel_size)
+    stride = _extract_scalar_param(module.stride, "stride", kernel_size)
+    padding = _extract_scalar_param(module.padding, "padding", 0)
     codegen.add_forward_call("MiCo_avgpool{dim}d_{dtype}", out, layer_name, input_names, 
-                             [kernel_size, module.stride, module.padding])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_module(torch.nn.MaxPool1d)
@@ -589,8 +626,10 @@ def handle_maxpool1d_module(codegen, n, out, module, input_names):
     layer_name = n.name
     codegen.add_uninitialized_tensor(layer_name, out)
     kernel_size = _extract_kernel_size(module.kernel_size)
+    stride = _extract_scalar_param(module.stride, "stride", kernel_size)
+    padding = _extract_scalar_param(module.padding, "padding", 0)
     codegen.add_forward_call("MiCo_maxpool{dim}d_{dtype}", out, layer_name, input_names, 
-                             [kernel_size, module.stride, module.padding])
+                             [kernel_size, stride, padding])
 
 
 @MiCoOpRegistry.register_module(torch.nn.AdaptiveAvgPool1d)

examples/mpq_train_bitnet.py

@@ -14,6 +14,7 @@
 argsparse.add_argument("--lr", type=float, default=0.001)
 argsparse.add_argument("-q", "--weight_quant", type=float, choices=[1,1.5,2], default=1)
 argsparse.add_argument("-aq", "--act_quant", type=int, choices=[4,8], default=8)
+argsparse.add_argument("--use-norm", action="store_true", default=False)
 argsparse.add_argument("--keep-last", action="store_true", default=False)
 argsparse.add_argument("--keep-first", action="store_true", default=False)
 argsparse.add_argument("--scheduler", type=str, default="none")
@@ -26,6 +27,7 @@
 scheduler = args.scheduler
 weight_quant = args.weight_quant
 act_quant = args.act_quant
+use_norm = args.use_norm
 keep_last = args.keep_last
 keep_first = args.keep_first
 
@@ -48,7 +50,7 @@
         qscheme[0][-1] = 8
         qscheme[1][-1] = 8
 
-    model.set_qscheme(qscheme, qat=True, use_norm=True)
+    model.set_qscheme(qscheme, qat=True, use_norm=use_norm)
     print("Model Param Size:", sum(p.numel() for p in model.parameters()))
     # Detect if there is a full precision checkpoint
     # if os.path.exists(f"output/ckpt/{model_name}.pth"):
@@ -75,7 +77,7 @@
     torch.save(model.state_dict(), f"output/ckpt/{model_name}_bitnet.pth")
     print("Model Train Results: ", res)
 
-    model.set_qscheme(qscheme, qat=True, use_norm=True)
+    model.set_qscheme(qscheme, qat=True, use_norm=use_norm)
 
     res = model.test(test_loader)
 

models/KWSTransformer.py

@@ -0,0 +1,113 @@
+import torch
+from torch import nn
+
+from MiCoModel import MiCoModel
+from models.CCT import TransformerClassifier
+
+
+class MFCCTokenizer(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 1,
+        embedding_dim: int = 96,
+        n_conv_layers: int = 2,
+        kernel_size=(3, 3),
+        stride=(1, 1),
+        padding=(1, 1),
+        pooling_kernel_size=(2, 2),
+        pooling_stride=(2, 2),
+        pooling_padding=(0, 0),
+    ):
+        super().__init__()
+
+        channels = [in_channels] + [embedding_dim] * n_conv_layers
+        layers = []
+        for i in range(n_conv_layers):
+            layers.extend(
+                [
+                    nn.Conv2d(
+                        channels[i],
+                        channels[i + 1],
+                        kernel_size=kernel_size,
+                        stride=stride,
+                        padding=padding,
+                        bias=False,
+                    ),
+                    nn.BatchNorm2d(channels[i + 1]),
+                    nn.ReLU(),
+                    nn.MaxPool2d(
+                        kernel_size=pooling_kernel_size,
+                        stride=pooling_stride,
+                        padding=pooling_padding,
+                    ),
+                ]
+            )
+
+        self.layers = nn.Sequential(*layers)
+        self.flatten = nn.Flatten(2, 3)
+
+    def sequence_length(self, input_size):
+        with torch.no_grad():
+            x = torch.zeros(1, 1, input_size[0], input_size[1])
+            return self.forward(x).shape[1]
+
+    def forward(self, x):
+        x = self.layers(x)
+        return self.flatten(x).transpose(1, 2)
+
+
+class KWSTransformer(MiCoModel):
+    def __init__(
+        self,
+        n_classes: int = 35,
+        input_size=(64, 81),
+        embedding_dim: int = 96,
+        n_conv_layers: int = 2,
+        num_layers: int = 2,
+        num_heads: int = 4,
+        mlp_ratio: float = 2.0,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        stochastic_depth: float = 0.0,
+        positional_embedding: str = "learnable",
+    ):
+        super().__init__()
+        self.default_dataset = "SPEECHCOMMANDS_2D"
+        self.input_size = tuple(input_size)
+
+        self.tokenizer = MFCCTokenizer(
+            in_channels=1,
+            embedding_dim=embedding_dim,
+            n_conv_layers=n_conv_layers,
+        )
+        self.classifier = TransformerClassifier(
+            sequence_length=self.tokenizer.sequence_length(self.input_size),
+            embedding_dim=embedding_dim,
+            seq_pool=True,
+            dropout=dropout,
+            attention_dropout=attention_dropout,
+            stochastic_depth=stochastic_depth,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            num_classes=n_classes,
+            positional_embedding=positional_embedding,
+        )
+        self.n_layers = len(self.get_qlayers())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.tokenizer(x)
+        return self.classifier(x)
+
+
+def tiny_kws_transformer(n_classes: int = 35):
+    return KWSTransformer(
+        n_classes=n_classes,
+        embedding_dim=64,
+        n_conv_layers=2,
+        num_layers=2,
+        num_heads=4,
+        mlp_ratio=2.0,
+        dropout=0.1,
+        attention_dropout=0.1,
+    )