diff --git a/MiCoONNXGen.py b/MiCoONNXGen.py new file mode 100644 index 0000000..778e2b4 --- /dev/null +++ b/MiCoONNXGen.py @@ -0,0 +1,267 @@ +import os +import json +import logging +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.nn +import torch.fx + +import onnx + +from MiCoQLayers import BitQLayer, BitConv2d, BitConv1d, BitLinear +from MiCoCodeGen import MiCoTrace + + +class MiCoONNXGen(torch.fx.Interpreter): + """ + ONNX exporter for mixed-precision quantized MiCo models. + + This class exports a PyTorch model to the ONNX format and attaches + per-layer bitwidth metadata (weight and activation quantization types) + so that downstream tools can reconstruct the mixed-precision configuration. + + Usage:: + + from MiCoONNXGen import MiCoONNXGen + from models import LeNet + from MiCoUtils import fuse_model + + model = LeNet(1) + model.set_qscheme([[8, 6, 6, 4, 4], [8, 8, 8, 8, 8]]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + exporter.export("output", "lenet_mnist", torch.randn(1, 1, 28, 28)) + """ + + def __init__(self, model: torch.nn.Module, log_level: int = logging.INFO): + graph, gm = MiCoONNXGen._extract_graph_module(model) + super().__init__(gm) + + self.model = model + self.graph = graph + self.gm = gm + self.logger = logging.getLogger("MiCoONNXGen") + self.logger.setLevel(log_level) + + # ------------------------------------------------------------------ + # Graph extraction (reuses MiCoTrace from MiCoCodeGen) + # ------------------------------------------------------------------ + @staticmethod + def _extract_graph_module(model: torch.nn.Module) -> Tuple[torch.fx.Graph, torch.fx.GraphModule]: + """Trace the model and return (graph, graph_module).""" + graph = MiCoTrace().trace(model) + graph.lint() + gm = torch.fx.GraphModule(model, graph) + return graph, gm + + # ------------------------------------------------------------------ + # Collect per-layer quantization metadata + # ------------------------------------------------------------------ + def _collect_bitwidth_info(self) -> Dict[str, Dict[str, Any]]: + """ + Walk the FX graph and collect bitwidth information for every + quantized layer (``BitQLayer`` subclasses). + + Returns: + A dictionary mapping layer names to their quantization + parameters, e.g.:: + + { + "layers.0": {"weight_bitwidth": 8, "activation_bitwidth": 8, "layer_type": "Conv2d"}, + ... + } + """ + info: Dict[str, Dict[str, Any]] = {} + + for node in self.graph.nodes: + if node.op != "call_module": + continue + + module = self._get_module(node.target) + if not isinstance(module, BitQLayer): + continue + + entry: Dict[str, Any] = { + "weight_bitwidth": int(module.qtype), + "activation_bitwidth": int(module.act_q), + } + + if isinstance(module, BitLinear): + entry["layer_type"] = "Linear" + elif isinstance(module, BitConv2d): + entry["layer_type"] = "Conv2d" + elif isinstance(module, BitConv1d): + entry["layer_type"] = "Conv1d" + else: + entry["layer_type"] = type(module).__name__ + + info[node.target] = entry + + return info + + def _get_module(self, target: str) -> torch.nn.Module: + """Resolve a dotted target path to the actual sub-module.""" + parts = target.split(".") + mod = self.model + for part in parts: + mod = getattr(mod, part) + return mod + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + def export( + self, + output_directory: str, + model_name: str, + example_input: torch.Tensor, + *, + opset_version: int = 18, + input_names: Optional[List[str]] = None, + output_names: Optional[List[str]] = None, + ) -> str: + """ + Export the model to ONNX with per-layer bitwidth metadata. + + Args: + output_directory: Directory where the ``.onnx`` file will be + written. + model_name: Base name for the output file (without extension). + example_input: A representative input tensor used by + ``torch.onnx.export`` for tracing. + opset_version: ONNX opset version (default 18). + input_names: Optional list of input names for the ONNX graph. + output_names: Optional list of output names for the ONNX graph. + + Returns: + The path to the written ``.onnx`` file. + """ + os.makedirs(output_directory, exist_ok=True) + onnx_path = os.path.join(output_directory, f"{model_name}.onnx") + + if input_names is None: + input_names = ["input"] + if output_names is None: + output_names = ["output"] + + # 1. Collect bitwidth information from the traced graph + bitwidth_info = self._collect_bitwidth_info() + + # 2. Export the model to ONNX via PyTorch + self.model.eval() + torch.onnx.export( + self.model, + example_input, + onnx_path, + opset_version=opset_version, + input_names=input_names, + output_names=output_names, + ) + + # 3. Re-load, attach metadata, and save + onnx_model = onnx.load(onnx_path) + + # Add whole-model metadata with the full bitwidth map + bitwidth_json = json.dumps(bitwidth_info) + onnx_model.metadata_props.append( + onnx.StringStringEntryProto( + key="mico_bitwidth_info", value=bitwidth_json + ) + ) + + # Also add per-node metadata as ONNX node attributes where possible + self._annotate_onnx_nodes(onnx_model, bitwidth_info) + + onnx.save(onnx_model, onnx_path) + + self.logger.info("ONNX model exported to %s", onnx_path) + self.logger.info( + "Per-layer bitwidth info (%d layers):\n%s", + len(bitwidth_info), + json.dumps(bitwidth_info, indent=2), + ) + + return onnx_path + + # ------------------------------------------------------------------ + # Helper: annotate ONNX graph nodes + # ------------------------------------------------------------------ + @staticmethod + def _annotate_onnx_nodes( + onnx_model: "onnx.ModelProto", + bitwidth_info: Dict[str, Dict[str, Any]], + ) -> None: + """ + Walk the ONNX graph and add ``weight_bitwidth`` / + ``activation_bitwidth`` doc_string annotations to nodes whose + names match the collected bitwidth info. + + Because ``torch.onnx.export`` flattens module hierarchy and may + rename nodes, this uses a best-effort name-matching heuristic: + a graph node is annotated if any key in *bitwidth_info* appears + as a substring of the node's name or output name. + """ + for node in onnx_model.graph.node: + matched_key = MiCoONNXGen._match_node_to_layer(node, bitwidth_info) + if matched_key is None: + continue + + entry = bitwidth_info[matched_key] + annotation = json.dumps( + { + "mico_layer": matched_key, + "weight_bitwidth": entry["weight_bitwidth"], + "activation_bitwidth": entry["activation_bitwidth"], + "layer_type": entry.get("layer_type", ""), + } + ) + node.doc_string = annotation + + @staticmethod + def _match_node_to_layer( + node: "onnx.NodeProto", + bitwidth_info: Dict[str, Dict[str, Any]], + ) -> Optional[str]: + """Return the bitwidth_info key that best matches *node*, or ``None``.""" + # Build a set of candidate strings from the node + candidates = [node.name] + list(node.output) + + for key in bitwidth_info: + # Normalize key: replace dots with underscores / slashes + normalized_variants = [ + key, + key.replace(".", "_"), + key.replace(".", "/"), + ] + for candidate in candidates: + if not candidate: + continue + for variant in normalized_variants: + if variant in candidate: + return key + return None + + # ------------------------------------------------------------------ + # Convenience: load and inspect metadata + # ------------------------------------------------------------------ + @staticmethod + def load_bitwidth_info(onnx_path: str) -> Dict[str, Dict[str, Any]]: + """ + Load an ONNX model and return the per-layer bitwidth info that + was embedded by :meth:`export`. + + Args: + onnx_path: Path to the ``.onnx`` file. + + Returns: + The bitwidth info dictionary, or an empty dict if no + metadata was found. + """ + onnx_model = onnx.load(onnx_path) + for prop in onnx_model.metadata_props: + if prop.key == "mico_bitwidth_info": + return json.loads(prop.value) + return {} diff --git a/tests/test_micoonnxgen.py b/tests/test_micoonnxgen.py new file mode 100644 index 0000000..cf0a3ed --- /dev/null +++ b/tests/test_micoonnxgen.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +""" +Test suite for MiCoONNXGen class. +Tests ONNX export functionality including initialization, +per-layer bitwidth metadata export, and model re-loading. +""" + +import torch +import torch.nn as nn +import sys +import os +import tempfile +import shutil +import json +import unittest + +# Add the parent directory to the path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from models import MLP, LeNet +from MiCoONNXGen import MiCoONNXGen +from MiCoUtils import fuse_model + + +class TestMiCoONNXGenInit(unittest.TestCase): + """Test MiCoONNXGen initialization.""" + + def test_init_mlp(self): + """Test initialization with an MLP model.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + self.assertIsNotNone(exporter) + self.assertIsNotNone(exporter.model) + self.assertIsNotNone(exporter.graph) + self.assertIsNotNone(exporter.gm) + + def test_init_lenet(self): + """Test initialization with a LeNet model.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + self.assertIsNotNone(exporter) + + +class TestMiCoONNXGenBitwidthCollection(unittest.TestCase): + """Test bitwidth metadata collection.""" + + def test_collect_uniform_bitwidths(self): + """Test collecting uniform (all-8-bit) bitwidths.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + info = exporter._collect_bitwidth_info() + + self.assertGreater(len(info), 0) + for layer_name, entry in info.items(): + self.assertEqual(entry["weight_bitwidth"], 8) + self.assertEqual(entry["activation_bitwidth"], 8) + self.assertIn("layer_type", entry) + + def test_collect_mixed_bitwidths(self): + """Test collecting mixed-precision bitwidths.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [4, 8] + activation_q = [8, 4] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + info = exporter._collect_bitwidth_info() + + self.assertEqual(len(info), 2) + bitwidths = list(info.values()) + self.assertEqual(bitwidths[0]["weight_bitwidth"], 4) + self.assertEqual(bitwidths[0]["activation_bitwidth"], 8) + self.assertEqual(bitwidths[1]["weight_bitwidth"], 8) + self.assertEqual(bitwidths[1]["activation_bitwidth"], 4) + + def test_collect_lenet_bitwidths(self): + """Test collecting bitwidths for LeNet (Conv2d + Linear layers).""" + model = LeNet(1) + weight_q = [8, 6, 6, 4, 4] + activation_q = [8, 8, 8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + info = exporter._collect_bitwidth_info() + + self.assertEqual(len(info), 5) + wbits = [e["weight_bitwidth"] for e in info.values()] + self.assertEqual(wbits, [8, 6, 6, 4, 4]) + + def test_layer_types_lenet(self): + """Test that layer types are correctly identified for LeNet.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + info = exporter._collect_bitwidth_info() + + types = [e["layer_type"] for e in info.values()] + # LeNet has 2 Conv2d layers followed by 3 Linear layers + self.assertEqual(types.count("Conv2d"), 2) + self.assertEqual(types.count("Linear"), 3) + + +class TestMiCoONNXGenExport(unittest.TestCase): + """Test ONNX export and metadata persistence.""" + + def setUp(self): + torch.manual_seed(42) + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + if os.path.exists(self.temp_dir): + shutil.rmtree(self.temp_dir) + + def test_export_mlp_creates_file(self): + """Test that export creates an .onnx file for MLP.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export(self.temp_dir, "test_mlp", torch.randn(1, 32)) + + self.assertTrue(os.path.exists(path)) + self.assertTrue(path.endswith(".onnx")) + + def test_export_lenet_creates_file(self): + """Test that export creates an .onnx file for LeNet.""" + model = LeNet(1) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export(self.temp_dir, "test_lenet", torch.randn(1, 1, 28, 28)) + + self.assertTrue(os.path.exists(path)) + + def test_export_metadata_roundtrip(self): + """Test that bitwidth metadata survives export -> load.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [4, 8] + activation_q = [8, 4] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export(self.temp_dir, "test_roundtrip", torch.randn(1, 32)) + + # Reload metadata + loaded_info = MiCoONNXGen.load_bitwidth_info(path) + + self.assertEqual(len(loaded_info), 2) + values = list(loaded_info.values()) + self.assertEqual(values[0]["weight_bitwidth"], 4) + self.assertEqual(values[0]["activation_bitwidth"], 8) + self.assertEqual(values[1]["weight_bitwidth"], 8) + self.assertEqual(values[1]["activation_bitwidth"], 4) + + def test_export_lenet_mixed_precision_roundtrip(self): + """Test mixed-precision metadata roundtrip for LeNet.""" + model = LeNet(1) + weight_q = [8, 6, 6, 4, 4] + activation_q = [8, 8, 8, 8, 8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export(self.temp_dir, "test_lenet_mpq", torch.randn(1, 1, 28, 28)) + + loaded_info = MiCoONNXGen.load_bitwidth_info(path) + + self.assertEqual(len(loaded_info), 5) + wbits = [e["weight_bitwidth"] for e in loaded_info.values()] + self.assertEqual(wbits, [8, 6, 6, 4, 4]) + + def test_export_custom_names(self): + """Test export with custom input/output names.""" + model = MLP(in_features=32, config={"Layers": [16, 10]}) + weight_q = [8] * model.n_layers + activation_q = [8] * model.n_layers + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export( + self.temp_dir, + "test_custom_names", + torch.randn(1, 32), + input_names=["my_input"], + output_names=["my_output"], + ) + + import onnx + onnx_model = onnx.load(path) + self.assertEqual(onnx_model.graph.input[0].name, "my_input") + self.assertEqual(onnx_model.graph.output[0].name, "my_output") + + def test_export_creates_directory(self): + """Test that export creates the output directory if needed.""" + nested_dir = os.path.join(self.temp_dir, "a", "b", "c") + model = MLP(in_features=16, config={"Layers": [8]}) + weight_q = [8] + activation_q = [8] + model.set_qscheme([weight_q, activation_q]) + model = fuse_model(model) + model.eval() + + exporter = MiCoONNXGen(model) + path = exporter.export(nested_dir, "nested_model", torch.randn(1, 16)) + + self.assertTrue(os.path.exists(path)) + + def test_load_bitwidth_info_no_metadata(self): + """Test load_bitwidth_info on a plain ONNX file without MiCo metadata.""" + # Export a plain model without MiCo metadata + model = nn.Sequential(nn.Linear(8, 4), nn.ReLU(), nn.Linear(4, 2)) + model.eval() + path = os.path.join(self.temp_dir, "plain.onnx") + torch.onnx.export(model, torch.randn(1, 8), path, opset_version=17) + + loaded_info = MiCoONNXGen.load_bitwidth_info(path) + self.assertEqual(loaded_info, {}) + + +if __name__ == "__main__": + unittest.main()