Commit 864b785
Fix BatchNorm fusion producing invalid ONNX when Conv nodes share weight initializers (#2883)
When a model reuses `Conv2d`+`BatchNorm2d` blocks (same weights called
twice), the BatchNorm fusion rewrite creates a new initializer with the
same name as the shared weight, overwriting it in the graph's
initializer dict. This sets `_is_initializer = False` on the original
value still referenced by the second Conv node, producing an invalid
model:
```
InvalidArgument: Node input 'conv.weight_1' is not a graph input, initializer, or output of a previous node.
```
Reproducer:
```python
class MWE(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 16, kernel_size=3, padding=1)
self.bn = nn.BatchNorm2d(16)
def forward(self, x):
f_1 = self.bn(self.conv(x[:, 0:1]))
f_2 = self.bn(self.conv(x[:, 1:2]))
return f_1 + f_2
torch.onnx.export(MWE().eval(), ..., optimize=True) # Invalid model
```
### Changes
- **`_fuse_batchnorm.py`**: Added a check in
`_FuseBatchNormBase.check()` that rejects the fusion when the inbound
node's weight/bias initializers are used by nodes outside the matched
pattern. This prevents the overwrite of shared initializers.
- **`_fuse_batchnorm_test.py`**: Added regression test with two Conv+BN
pairs sharing the same weight and bias initializers.
---------
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>1 parent 12234f8 commit 864b785
File tree
2 files changed
+67
-0
lines changed- onnxscript/rewriter/rules/common
2 files changed
+67
-0
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
103 | 103 | | |
104 | 104 | | |
105 | 105 | | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
106 | 123 | | |
107 | 124 | | |
108 | 125 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
311 | 311 | | |
312 | 312 | | |
313 | 313 | | |
| 314 | + | |
| 315 | + | |
| 316 | + | |
| 317 | + | |
| 318 | + | |
| 319 | + | |
| 320 | + | |
| 321 | + | |
| 322 | + | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
| 329 | + | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
| 335 | + | |
| 336 | + | |
| 337 | + | |
| 338 | + | |
| 339 | + | |
| 340 | + | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
| 359 | + | |
| 360 | + | |
| 361 | + | |
| 362 | + | |
| 363 | + | |
314 | 364 | | |
315 | 365 | | |
316 | 366 | | |
0 commit comments