Support fuse bn into ConvTranspose. - #316
Open
take-cheeze wants to merge 5 commits into
Open
Conversation
Signed-off-by: wenyuchi.wyc <wenyuchi.wyc@alibaba-inc.com>
Signed-off-by: Takeshi Watanabe <take-cheeze@users.noreply.github.com>
Extend the fuse_bn_into_conv pass to fold BatchNormalization into a preceding ConvTranspose, and fix the channel-axis bug from PR onnx#106 that crashed on real models. ConvTranspose weight is laid out as (in_channels, out_channels, kH, kW), transposed relative to Conv's (out_channels, in_channels, kH, kW). The BatchNormalization channel count matches the output channels, so the per-channel scale must be checked and broadcast against axis 1 for ConvTranspose (axis 0 for Conv). Using axis 0 unconditionally triggered the reported `conv_W.sizes()[0] == C` assertion failure whenever in_channels != out_channels. Mismatched shapes (e.g. grouped ConvTranspose) now skip the fusion instead of asserting. The added test uses distinct in/out channel counts so it actually exercises the corrected axis. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LuHDzT8V8x1FbThhpWG6ad Signed-off-by: take-cheeze <takechi101010@gmail.com>
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
3 participants
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.
Based and fixes #106
Summary
This PR extends the
fuse_bn_into_convoptimization pass to support fusing BatchNormalization nodes into ConvTranspose operations, in addition to the existing Conv support.Key Changes
patternMatchPredicateto recognize bothConvandConvTransposeoperations followed byBatchNormalizationmodify_convto accept anis_convparameter that accounts for the different weight tensor layouts:(out_channels, in_channels/group, kH, kW)with output channels on axis 0(in_channels, out_channels/group, kH, kW)with output channels on axis 1test_fuse_bn_into_conv_transpose_simpleto verify the fusion works correctly with distinct input/output channel countsImplementation Details
conv_W.sizes()[0] == C) that occurred when attempting to fuse BatchNormalization into ConvTranspose with mismatched input/output channel countsstd::iotaanderaseto generate the correct unsqueeze dimensions dynamically based on the operation type#include <numeric>forstd::iotasupporthttps://claude.ai/code/session_01Fquk8PCqZ2T26G6CwLAuDn