Open
Add MiCoONNXGen ONNX exporter with per-layer bitwidth metadata#37
Conversation
Co-authored-by: Jzjerry <20167827+Jzjerry@users.noreply.github.com>
Co-authored-by: Jzjerry <20167827+Jzjerry@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Add ONNX exporter for mixed-precision networks
Add MiCoONNXGen ONNX exporter with per-layer bitwidth metadata
Mar 12, 2026
Contributor
|
@sourcery-ai review |
Reviewer's GuideIntroduces a MiCoONNXGen ONNX exporter that traces MiCo mixed-precision models with FX, collects per-layer BitQLayer bitwidth metadata, embeds it into the exported ONNX model (both as global metadata and per-node annotations), and adds a comprehensive unit test suite validating graph extraction, metadata collection, export behavior, and metadata roundtrip. Sequence diagram for ONNX export with embedded bitwidth metadatasequenceDiagram
actor User
participant MiCoONNXGen
participant MiCoTrace
participant TorchONNX as Torch_ONNX
participant ONNX as ONNX_Runtime
User->>MiCoONNXGen: MiCoONNXGen(model)
activate MiCoONNXGen
MiCoONNXGen->>MiCoONNXGen: _extract_graph_module(model)
MiCoONNXGen->>MiCoTrace: trace(model)
activate MiCoTrace
MiCoTrace-->>MiCoONNXGen: torch_fx_Graph
deactivate MiCoTrace
MiCoONNXGen->>MiCoONNXGen: create torch_fx_GraphModule
MiCoONNXGen-->>User: exporter instance
deactivate MiCoONNXGen
User->>MiCoONNXGen: export(output_directory, model_name, example_input)
activate MiCoONNXGen
MiCoONNXGen->>MiCoONNXGen: _collect_bitwidth_info()
MiCoONNXGen-->>MiCoONNXGen: bitwidth_info Dict
MiCoONNXGen->>TorchONNX: export(model, example_input, onnx_path, opset_version, input_names, output_names)
activate TorchONNX
TorchONNX-->>MiCoONNXGen: write ONNX file
deactivate TorchONNX
MiCoONNXGen->>ONNX: load(onnx_path)
activate ONNX
ONNX-->>MiCoONNXGen: onnx_model
MiCoONNXGen->>MiCoONNXGen: attach metadata_props[mico_bitwidth_info]
MiCoONNXGen->>MiCoONNXGen: _annotate_onnx_nodes(onnx_model, bitwidth_info)
MiCoONNXGen->>ONNX: save(onnx_model, onnx_path)
ONNX-->>MiCoONNXGen: onnx_path
deactivate ONNX
MiCoONNXGen-->>User: onnx_path
deactivate MiCoONNXGen
User->>MiCoONNXGen: load_bitwidth_info(onnx_path)
activate MiCoONNXGen
MiCoONNXGen->>ONNX: load(onnx_path)
activate ONNX
ONNX-->>MiCoONNXGen: onnx_model
deactivate ONNX
MiCoONNXGen->>MiCoONNXGen: scan metadata_props for mico_bitwidth_info
MiCoONNXGen-->>User: bitwidth_info Dict
deactivate MiCoONNXGen
Class diagram for MiCoONNXGen ONNX exporter with per-layer bitwidth metadataclassDiagram
class torch_fx_Interpreter {
}
class MiCoTrace {
+trace(model) torch_fx_Graph
}
class BitQLayer {
+qtype int
+act_q int
}
class BitConv2d {
}
class BitConv1d {
}
class BitLinear {
}
class MiCoONNXGen {
-model torch_nn_Module
-graph torch_fx_Graph
-gm torch_fx_GraphModule
-logger logging_Logger
+MiCoONNXGen(model, log_level)
+export(output_directory, model_name, example_input, opset_version, input_names, output_names) str
+load_bitwidth_info(onnx_path) Dict
-_extract_graph_module(model) torch_fx_Graph torch_fx_GraphModule
-_collect_bitwidth_info() Dict
-_get_module(target) torch_nn_Module
-_annotate_onnx_nodes(onnx_model, bitwidth_info) void
-_match_node_to_layer(node, bitwidth_info) str
}
MiCoONNXGen --|> torch_fx_Interpreter
BitConv2d --|> BitQLayer
BitConv1d --|> BitQLayer
BitLinear --|> BitQLayer
MiCoONNXGen ..> MiCoTrace : uses for tracing
MiCoONNXGen ..> BitQLayer : collects bitwidth
MiCoONNXGen ..> BitConv2d
MiCoONNXGen ..> BitConv1d
MiCoONNXGen ..> BitLinear
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The best-effort name matching in
_match_node_to_layeris purely substring-based and may mis-associate layers in larger graphs; consider tightening the heuristic (e.g., anchoring on full-name variants or verifying op types) or allowing the caller to inject a mapping when needed. - In
_annotate_onnx_nodes, settingnode.doc_stringdirectly overwrites any existing documentation; if preserving human-readable docs is important, consider appending the JSON annotation or storing it as explicit ONNX node attributes instead. load_bitwidth_infoassumes well-formed JSON inmico_bitwidth_infoand will raise if the metadata is corrupted or manually edited; adding a small try/except with a fallback to{}(and perhaps a log warning) would make this more robust.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The best-effort name matching in `_match_node_to_layer` is purely substring-based and may mis-associate layers in larger graphs; consider tightening the heuristic (e.g., anchoring on full-name variants or verifying op types) or allowing the caller to inject a mapping when needed.
- In `_annotate_onnx_nodes`, setting `node.doc_string` directly overwrites any existing documentation; if preserving human-readable docs is important, consider appending the JSON annotation or storing it as explicit ONNX node attributes instead.
- `load_bitwidth_info` assumes well-formed JSON in `mico_bitwidth_info` and will raise if the metadata is corrupted or manually edited; adding a small try/except with a fallback to `{}` (and perhaps a log warning) would make this more robust.
## Individual Comments
### Comment 1
<location path="MiCoONNXGen.py" line_range="213-220" />
<code_context>
+ "layer_type": entry.get("layer_type", ""),
+ }
+ )
+ node.doc_string = annotation
+
+ @staticmethod
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Avoid overwriting existing ONNX node doc_strings when annotating bitwidth metadata.
This always overwrites `node.doc_string`, potentially discarding metadata added by other tools or passes. Instead, preserve the existing value (for example by appending your JSON in a structured way or storing it in a separate node attribute) so prior annotations remain available.
```suggestion
annotation = json.dumps(
{
"mico_layer": matched_key,
"weight_bitwidth": entry["weight_bitwidth"],
"activation_bitwidth": entry["activation_bitwidth"],
"layer_type": entry.get("layer_type", ""),
}
)
# Preserve existing doc_string content by embedding MiCo metadata alongside it
existing_doc = getattr(node, "doc_string", None)
if existing_doc:
try:
existing_json = json.loads(existing_doc)
except Exception:
# Existing doc_string is not JSON: wrap both pieces in a new JSON object
combined = {
"previous_doc_string": existing_doc,
"mico_bitwidth": json.loads(annotation),
}
else:
if isinstance(existing_json, dict):
# Extend existing JSON doc with MiCo bitwidth metadata in a dedicated field
combined = dict(existing_json)
combined["mico_bitwidth"] = json.loads(annotation)
else:
# Existing JSON is not a dict: wrap to avoid losing structure
combined = {
"previous_doc_string": existing_doc,
"mico_bitwidth": json.loads(annotation),
}
node.doc_string = json.dumps(combined)
else:
# No existing doc_string: use MiCo metadata as-is
node.doc_string = annotation
```
</issue_to_address>
### Comment 2
<location path="MiCoONNXGen.py" line_range="264-267" />
<code_context>
+ 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 {}
</code_context>
<issue_to_address>
**suggestion:** load_bitwidth_info assumes the metadata value is always valid JSON.
If `mico_bitwidth_info` contains invalid or non-JSON content, `json.loads(prop.value)` will raise and fail this call. Consider wrapping the parse in a try/except and returning `{}` (and/or logging a warning) when JSON decoding fails.
Suggested implementation:
```python
Returns:
The bitwidth info dictionary, or an empty dict if no
metadata was found or if the metadata content is invalid.
```
```python
onnx_model = onnx.load(onnx_path)
for prop in onnx_model.metadata_props:
if prop.key == "mico_bitwidth_info":
try:
return json.loads(prop.value)
except (TypeError, ValueError):
# Return empty dict if metadata content is not valid JSON
return {}
return {}
```
</issue_to_address>
### Comment 3
<location path="tests/test_micoonnxgen.py" line_range="69-78" />
<code_context>
+ 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):
</code_context>
<issue_to_address>
**suggestion (testing):** Avoid relying on dict iteration order in bitwidth assertions to prevent potential flakiness.
These assertions depend on `dict.values()` reflecting a specific layer order, which in turn depends on FX traversal and may change with model or tracing updates. To make the tests stable, index by layer name (e.g., `info["layers.0"]`, `info["layers.1"]`) or sort items by key before asserting. The same issue appears in `test_export_metadata_roundtrip` and `test_export_lenet_mixed_precision_roundtrip`.
Suggested implementation:
```python
exporter = MiCoONNXGen(model)
info = exporter._collect_bitwidth_info()
self.assertEqual(len(info), 2)
# Sort by layer name to avoid depending on dict insertion / FX traversal order
sorted_items = sorted(info.items(), key=lambda kv: kv[0])
first_layer = sorted_items[0][1]
second_layer = sorted_items[1][1]
self.assertEqual(first_layer["weight_bitwidth"], 4)
self.assertEqual(first_layer["activation_bitwidth"], 8)
self.assertEqual(second_layer["weight_bitwidth"], 8)
self.assertEqual(second_layer["activation_bitwidth"], 4)
```
The same non-deterministic reliance on `dict.values()` appears in `test_export_metadata_roundtrip` and `test_export_lenet_mixed_precision_roundtrip`. In those tests, replace usages of `list(info.values())[idx]` (or similar index-based access on `dict.values()`/`dict.items()` without ordering) with a deterministic access pattern, e.g.:
1. Either:
- Index by explicit layer names: `info["layers.0"]`, `info["layers.1"]`, etc., if those keys are stable and known in the test.
2. Or:
- Sort by key before indexing:
```python
sorted_items = sorted(info.items(), key=lambda kv: kv[0])
layer0 = sorted_items[0][1]
layer1 = sorted_items[1][1]
```
Apply whichever pattern matches the expectations of each test and the known structure of the `info` mapping.
</issue_to_address>
### Comment 4
<location path="tests/test_micoonnxgen.py" line_range="130-139" />
<code_context>
+class TestMiCoONNXGenExport(unittest.TestCase):
</code_context>
<issue_to_address>
**suggestion (testing):** Add tests to cover ONNX node-level annotations (`doc_string`) created by `_annotate_onnx_nodes` and `_match_node_to_layer`.
The per-layer bitwidth metadata is also attached to individual ONNX nodes via `doc_string` (via `_annotate_onnx_nodes` and `_match_node_to_layer`), but the current tests don’t validate this. Please add focused tests that (1) export a small MiCo model, inspect the ONNX graph, and assert that some `graph.node` entries have non-empty `doc_string` values containing valid JSON with `mico_layer`, `weight_bitwidth`, `activation_bitwidth`, and `layer_type`; and (2) verify the name-matching heuristic for common patterns (e.g., module names with dots vs. node names with underscores) and that non-`BitQLayer` nodes remain unannotated. This will ensure the per-node metadata and matching logic are exercised and protected against regressions.
Suggested implementation:
```python
import tempfile
import json
import onnx
```
```python
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 _export_small_mico_model_to_onnx(self, model_name: str = "lenet") -> onnx.ModelProto:
"""Helper that builds a small MiCo model, exports it to ONNX, and returns the loaded ONNX model.
`model_name` is left as an argument to make it easy to plug in different
test models (e.g., LeNet, toy CNN) depending on what the test utilities provide.
"""
# Build a small MiCo model – adjust this to match the existing factory/helper.
# For example, if there is already a `build_lenet_model()` or similar in this
# test file or a shared utilities module, call that instead.
mico_model = build_lenet_model() if model_name == "lenet" else build_mico_model(model_name)
exporter = MiCoONNXGen(mico_model)
onnx_path = os.path.join(self.temp_dir, f"{model_name}.onnx")
# Adjust this call to match the actual ONNX export API exposed by MiCoONNXGen.
exporter.export_onnx(onnx_path)
return onnx.load(onnx_path)
def test_onnx_nodes_have_bitwidth_docstring_metadata(self):
"""Export a small MiCo model and verify ONNX nodes carry per-layer bitwidth metadata in doc_string."""
onnx_model = self._export_small_mico_model_to_onnx("lenet")
annotated_nodes = []
for node in onnx_model.graph.node:
if not node.doc_string:
continue
try:
meta = json.loads(node.doc_string)
except json.JSONDecodeError:
self.fail(f"ONNX node doc_string is not valid JSON: {node.doc_string!r}")
# Basic shape of the metadata
self.assertIn("mico_layer", meta)
self.assertIn("weight_bitwidth", meta)
self.assertIn("activation_bitwidth", meta)
self.assertIn("layer_type", meta)
# Values should be non-trivial / non-empty
self.assertIsInstance(meta["mico_layer"], str)
self.assertNotEqual(meta["mico_layer"], "")
self.assertIsInstance(meta["layer_type"], str)
self.assertNotEqual(meta["layer_type"], "")
self.assertIsInstance(meta["weight_bitwidth"], int)
self.assertGreater(meta["weight_bitwidth"], 0)
self.assertIsInstance(meta["activation_bitwidth"], int)
self.assertGreater(meta["activation_bitwidth"], 0)
annotated_nodes.append(node)
# We expect at least some nodes (those corresponding to BitQLayer instances)
# to carry doc_string annotations.
self.assertGreater(
len(annotated_nodes),
0,
"Expected at least one ONNX node to have MiCo bitwidth metadata in doc_string",
)
def test_onnx_node_name_matching_and_unannotated_non_bitq_layers(self):
"""Verify name-matching heuristic for node annotations and that non-BitQLayer nodes remain unannotated.
This exercises `_match_node_to_layer` by checking that:
* Nodes corresponding to BitQLayer modules (with dotted names) get matched and annotated.
* Nodes that do not correspond to any BitQLayer remain without doc_string metadata.
"""
onnx_model = self._export_small_mico_model_to_onnx("lenet")
annotated_layer_names = set()
unannotated_nodes = []
for node in onnx_model.graph.node:
if node.doc_string:
try:
meta = json.loads(node.doc_string)
except json.JSONDecodeError:
self.fail(f"ONNX node doc_string is not valid JSON: {node.doc_string!r}")
mico_layer_name = meta.get("mico_layer", "")
# Collect the mapped MiCo layer names for later assertions
if isinstance(mico_layer_name, str) and mico_layer_name:
annotated_layer_names.add(mico_layer_name)
# Ensure that layers that are annotated have corresponding basic fields
self.assertIn("layer_type", meta)
else:
unannotated_nodes.append(node)
# At least one annotated layer name should look like a typical PyTorch module path
# (containing dots), exercising the "dots vs underscores" matching heuristic.
has_dotted_layer_name = any("." in name for name in annotated_layer_names)
self.assertTrue(
has_dotted_layer_name,
"Expected at least one annotated node to map to a MiCo layer with a dotted module name",
)
# And we should still have some nodes that remain unannotated, representing
# nodes not corresponding to BitQLayer modules (e.g., shape/constant/graph utility nodes).
self.assertGreater(
len(unannotated_nodes),
0,
"Expected some ONNX nodes to remain unannotated when they do not map to BitQLayer modules",
)
```
1. Replace `build_lenet_model()` / `build_mico_model(model_name)` in `_export_small_mico_model_to_onnx` with the actual helper or factory used in your tests to create a small MiCo-quantized model (e.g., a LeNet variant with `BitQLayer` modules).
2. Adjust `exporter.export_onnx(onnx_path)` to match the real ONNX export method exposed by `MiCoONNXGen` (e.g., `export`, `export_onnx`, or similar), including any additional required arguments (like input shape or dummy input tensors).
3. If `onnx` and/or `json` are already imported at the top of `tests/test_micoonnxgen.py`, you can remove the duplicate imports from the first edit block.
4. If the test suite already provides a canonical way to reference a "small MiCo model" (e.g., a fixture or shared utility function), you may want to parameterize the tests or reuse that helper instead of hardcoding `"lenet"`.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+213
to
+220
| annotation = json.dumps( | ||
| { | ||
| "mico_layer": matched_key, | ||
| "weight_bitwidth": entry["weight_bitwidth"], | ||
| "activation_bitwidth": entry["activation_bitwidth"], | ||
| "layer_type": entry.get("layer_type", ""), | ||
| } | ||
| ) |
There was a problem hiding this comment.
suggestion (bug_risk): Avoid overwriting existing ONNX node doc_strings when annotating bitwidth metadata.
This always overwrites node.doc_string, potentially discarding metadata added by other tools or passes. Instead, preserve the existing value (for example by appending your JSON in a structured way or storing it in a separate node attribute) so prior annotations remain available.
Suggested change
| annotation = json.dumps( | |
| { | |
| "mico_layer": matched_key, | |
| "weight_bitwidth": entry["weight_bitwidth"], | |
| "activation_bitwidth": entry["activation_bitwidth"], | |
| "layer_type": entry.get("layer_type", ""), | |
| } | |
| ) | |
| annotation = json.dumps( | |
| { | |
| "mico_layer": matched_key, | |
| "weight_bitwidth": entry["weight_bitwidth"], | |
| "activation_bitwidth": entry["activation_bitwidth"], | |
| "layer_type": entry.get("layer_type", ""), | |
| } | |
| ) | |
| # Preserve existing doc_string content by embedding MiCo metadata alongside it | |
| existing_doc = getattr(node, "doc_string", None) | |
| if existing_doc: | |
| try: | |
| existing_json = json.loads(existing_doc) | |
| except Exception: | |
| # Existing doc_string is not JSON: wrap both pieces in a new JSON object | |
| combined = { | |
| "previous_doc_string": existing_doc, | |
| "mico_bitwidth": json.loads(annotation), | |
| } | |
| else: | |
| if isinstance(existing_json, dict): | |
| # Extend existing JSON doc with MiCo bitwidth metadata in a dedicated field | |
| combined = dict(existing_json) | |
| combined["mico_bitwidth"] = json.loads(annotation) | |
| else: | |
| # Existing JSON is not a dict: wrap to avoid losing structure | |
| combined = { | |
| "previous_doc_string": existing_doc, | |
| "mico_bitwidth": json.loads(annotation), | |
| } | |
| node.doc_string = json.dumps(combined) | |
| else: | |
| # No existing doc_string: use MiCo metadata as-is | |
| node.doc_string = annotation |
Comment on lines
+264
to
+267
| for prop in onnx_model.metadata_props: | ||
| if prop.key == "mico_bitwidth_info": | ||
| return json.loads(prop.value) | ||
| return {} |
There was a problem hiding this comment.
suggestion: load_bitwidth_info assumes the metadata value is always valid JSON.
If mico_bitwidth_info contains invalid or non-JSON content, json.loads(prop.value) will raise and fail this call. Consider wrapping the parse in a try/except and returning {} (and/or logging a warning) when JSON decoding fails.
Suggested implementation:
Returns:
The bitwidth info dictionary, or an empty dict if no
metadata was found or if the metadata content is invalid. onnx_model = onnx.load(onnx_path)
for prop in onnx_model.metadata_props:
if prop.key == "mico_bitwidth_info":
try:
return json.loads(prop.value)
except (TypeError, ValueError):
# Return empty dict if metadata content is not valid JSON
return {}
return {}
Comment on lines
+69
to
+78
| 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.""" |
There was a problem hiding this comment.
suggestion (testing): Avoid relying on dict iteration order in bitwidth assertions to prevent potential flakiness.
These assertions depend on dict.values() reflecting a specific layer order, which in turn depends on FX traversal and may change with model or tracing updates. To make the tests stable, index by layer name (e.g., info["layers.0"], info["layers.1"]) or sort items by key before asserting. The same issue appears in test_export_metadata_roundtrip and test_export_lenet_mixed_precision_roundtrip.
Suggested implementation:
exporter = MiCoONNXGen(model)
info = exporter._collect_bitwidth_info()
self.assertEqual(len(info), 2)
# Sort by layer name to avoid depending on dict insertion / FX traversal order
sorted_items = sorted(info.items(), key=lambda kv: kv[0])
first_layer = sorted_items[0][1]
second_layer = sorted_items[1][1]
self.assertEqual(first_layer["weight_bitwidth"], 4)
self.assertEqual(first_layer["activation_bitwidth"], 8)
self.assertEqual(second_layer["weight_bitwidth"], 8)
self.assertEqual(second_layer["activation_bitwidth"], 4)The same non-deterministic reliance on dict.values() appears in test_export_metadata_roundtrip and test_export_lenet_mixed_precision_roundtrip. In those tests, replace usages of list(info.values())[idx] (or similar index-based access on dict.values()/dict.items() without ordering) with a deterministic access pattern, e.g.:
- Either:
- Index by explicit layer names:
info["layers.0"],info["layers.1"], etc., if those keys are stable and known in the test.
- Index by explicit layer names:
- Or:
- Sort by key before indexing:
sorted_items = sorted(info.items(), key=lambda kv: kv[0]) layer0 = sorted_items[0][1] layer1 = sorted_items[1][1]
- Sort by key before indexing:
Apply whichever pattern matches the expectations of each test and the known structure of the info mapping.
Comment on lines
+130
to
+139
| 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) |
There was a problem hiding this comment.
suggestion (testing): Add tests to cover ONNX node-level annotations (doc_string) created by _annotate_onnx_nodes and _match_node_to_layer.
The per-layer bitwidth metadata is also attached to individual ONNX nodes via doc_string (via _annotate_onnx_nodes and _match_node_to_layer), but the current tests don’t validate this. Please add focused tests that (1) export a small MiCo model, inspect the ONNX graph, and assert that some graph.node entries have non-empty doc_string values containing valid JSON with mico_layer, weight_bitwidth, activation_bitwidth, and layer_type; and (2) verify the name-matching heuristic for common patterns (e.g., module names with dots vs. node names with underscores) and that non-BitQLayer nodes remain unannotated. This will ensure the per-node metadata and matching logic are exercised and protected against regressions.
Suggested implementation:
import tempfile
import json
import onnxclass 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 _export_small_mico_model_to_onnx(self, model_name: str = "lenet") -> onnx.ModelProto:
"""Helper that builds a small MiCo model, exports it to ONNX, and returns the loaded ONNX model.
`model_name` is left as an argument to make it easy to plug in different
test models (e.g., LeNet, toy CNN) depending on what the test utilities provide.
"""
# Build a small MiCo model – adjust this to match the existing factory/helper.
# For example, if there is already a `build_lenet_model()` or similar in this
# test file or a shared utilities module, call that instead.
mico_model = build_lenet_model() if model_name == "lenet" else build_mico_model(model_name)
exporter = MiCoONNXGen(mico_model)
onnx_path = os.path.join(self.temp_dir, f"{model_name}.onnx")
# Adjust this call to match the actual ONNX export API exposed by MiCoONNXGen.
exporter.export_onnx(onnx_path)
return onnx.load(onnx_path)
def test_onnx_nodes_have_bitwidth_docstring_metadata(self):
"""Export a small MiCo model and verify ONNX nodes carry per-layer bitwidth metadata in doc_string."""
onnx_model = self._export_small_mico_model_to_onnx("lenet")
annotated_nodes = []
for node in onnx_model.graph.node:
if not node.doc_string:
continue
try:
meta = json.loads(node.doc_string)
except json.JSONDecodeError:
self.fail(f"ONNX node doc_string is not valid JSON: {node.doc_string!r}")
# Basic shape of the metadata
self.assertIn("mico_layer", meta)
self.assertIn("weight_bitwidth", meta)
self.assertIn("activation_bitwidth", meta)
self.assertIn("layer_type", meta)
# Values should be non-trivial / non-empty
self.assertIsInstance(meta["mico_layer"], str)
self.assertNotEqual(meta["mico_layer"], "")
self.assertIsInstance(meta["layer_type"], str)
self.assertNotEqual(meta["layer_type"], "")
self.assertIsInstance(meta["weight_bitwidth"], int)
self.assertGreater(meta["weight_bitwidth"], 0)
self.assertIsInstance(meta["activation_bitwidth"], int)
self.assertGreater(meta["activation_bitwidth"], 0)
annotated_nodes.append(node)
# We expect at least some nodes (those corresponding to BitQLayer instances)
# to carry doc_string annotations.
self.assertGreater(
len(annotated_nodes),
0,
"Expected at least one ONNX node to have MiCo bitwidth metadata in doc_string",
)
def test_onnx_node_name_matching_and_unannotated_non_bitq_layers(self):
"""Verify name-matching heuristic for node annotations and that non-BitQLayer nodes remain unannotated.
This exercises `_match_node_to_layer` by checking that:
* Nodes corresponding to BitQLayer modules (with dotted names) get matched and annotated.
* Nodes that do not correspond to any BitQLayer remain without doc_string metadata.
"""
onnx_model = self._export_small_mico_model_to_onnx("lenet")
annotated_layer_names = set()
unannotated_nodes = []
for node in onnx_model.graph.node:
if node.doc_string:
try:
meta = json.loads(node.doc_string)
except json.JSONDecodeError:
self.fail(f"ONNX node doc_string is not valid JSON: {node.doc_string!r}")
mico_layer_name = meta.get("mico_layer", "")
# Collect the mapped MiCo layer names for later assertions
if isinstance(mico_layer_name, str) and mico_layer_name:
annotated_layer_names.add(mico_layer_name)
# Ensure that layers that are annotated have corresponding basic fields
self.assertIn("layer_type", meta)
else:
unannotated_nodes.append(node)
# At least one annotated layer name should look like a typical PyTorch module path
# (containing dots), exercising the "dots vs underscores" matching heuristic.
has_dotted_layer_name = any("." in name for name in annotated_layer_names)
self.assertTrue(
has_dotted_layer_name,
"Expected at least one annotated node to map to a MiCo layer with a dotted module name",
)
# And we should still have some nodes that remain unannotated, representing
# nodes not corresponding to BitQLayer modules (e.g., shape/constant/graph utility nodes).
self.assertGreater(
len(unannotated_nodes),
0,
"Expected some ONNX nodes to remain unannotated when they do not map to BitQLayer modules",
)- Replace
build_lenet_model()/build_mico_model(model_name)in_export_small_mico_model_to_onnxwith the actual helper or factory used in your tests to create a small MiCo-quantized model (e.g., a LeNet variant withBitQLayermodules). - Adjust
exporter.export_onnx(onnx_path)to match the real ONNX export method exposed byMiCoONNXGen(e.g.,export,export_onnx, or similar), including any additional required arguments (like input shape or dummy input tensors). - If
onnxand/orjsonare already imported at the top oftests/test_micoonnxgen.py, you can remove the duplicate imports from the first edit block. - If the test suite already provides a canonical way to reference a "small MiCo model" (e.g., a fixture or shared utility function), you may want to parameterize the tests or reuse that helper instead of hardcoding
"lenet".
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
New
MiCoONNXGenclass that exports mixed-precision quantized models to ONNX, embedding per-layer weight/activation bitwidth settings as model metadata.MiCoONNXGen.py: Extendstorch.fx.Interpreter(mirrorsMiCoCodeGenpattern), reusesMiCoTracefor graph extraction. Walks FX graph to collectBitQLayerbitwidth info, exports viatorch.onnx.export, then attaches bitwidth map as ONNXmetadata_props(key:mico_bitwidth_info) and annotates individual ONNX nodes viadoc_string.tests/test_micoonnxgen.py: 13 tests covering init, bitwidth collection (uniform/mixed/Conv2d+Linear), ONNX file creation, metadata roundtrip, custom names, and edge cases.💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.
Summary by Sourcery
Add an ONNX exporter for MiCo mixed-precision quantized models that embeds per-layer bitwidth metadata into exported models and provides utilities to reload this metadata.
New Features:
Tests: