Skip to content

Support fuse bn into ConvTranspose.#106

Open
YuchiWen wants to merge 4 commits into
onnx:mainfrom
YuchiWen:fuse_bn_into_conv_transpose
Open

Support fuse bn into ConvTranspose.#106
YuchiWen wants to merge 4 commits into
onnx:mainfrom
YuchiWen:fuse_bn_into_conv_transpose

Conversation

@YuchiWen

Copy link
Copy Markdown

No description provided.

@YuchiWen
YuchiWen force-pushed the fuse_bn_into_conv_transpose branch from 44658d1 to 6e5ac70 Compare November 22, 2022 04:01
@daquexian

Copy link
Copy Markdown
Member

Sorry for the late response. Could you please add some tests for the fusion? You can follow the conv-bn fusion https://github.com/onnx/optimizer/blob/master/onnxoptimizer/test/optimizer_test.py#L3024

@YuchiWen
YuchiWen force-pushed the fuse_bn_into_conv_transpose branch from 6e5ac70 to 5d4d388 Compare March 6, 2023 10:52
@YuchiWen

YuchiWen commented Mar 6, 2023

Copy link
Copy Markdown
Author

Sorry for the late response. Could you please add some tests for the fusion? You can follow the conv-bn fusion https://github.com/onnx/optimizer/blob/master/onnxoptimizer/test/optimizer_test.py#L3024

@daquexian Done, please review.

Signed-off-by: wenyuchi.wyc <wenyuchi.wyc@alibaba-inc.com>
@YuchiWen
YuchiWen force-pushed the fuse_bn_into_conv_transpose branch from 5d4d388 to ff1229e Compare March 6, 2023 11:31
@huangzhicong3

Copy link
Copy Markdown

Hello, i try to used this commit to fuse the bn layer and convtranspose layer in my model and find some bugs:
The error message is:
passes/fuse_bn_into_conv.h:71: modify_conv: Assertion conv_W.sizes().size() > 2 && conv_W.sizes()[0] == C failed.

From the doc of onnx website (https://onnx.ai/onnx/operators/onnx__ConvTranspose.html), the shape of weight array of convtranspose is (Cin, Cout, K, K), which is different to normal Conv layer (Cout, Cin, K, K).

@huangzhicong3

huangzhicong3 commented May 30, 2024

Copy link
Copy Markdown

Hi, i would like to share my codes for fusing convtranspose and bn. It has been tested on my own model. I hope it will help others who have the same issue.

import numpy as np
import onnx
import sclblonnx as so

model = onnx.load('../onnx/models/backbone_clean.onnx')

all_initializer = model.graph.initializer
all_node = model.graph.node
ConvTranspose_list = []
BatchNormalization_list = []
for i, node in enumerate(all_node):
    # search convtranspose and batchnormalization
    if node.op_type == "ConvTranspose":
        # print(i, node.name, node.op_type,  node.input, node.output)
        ConvTranspose_list.append(node)
    if node.op_type == "BatchNormalization":
        # print(i, node.name, node.op_type,  node.input, node.output)
        BatchNormalization_list.append(node)

valid_ConvTranspose_list = []
for node in ConvTranspose_list:
    output = node.output
    for bn_node in BatchNormalization_list:
        bn_inputs = bn_node.input
        if output[0] in bn_inputs:
            valid_ConvTranspose_list.append({"conv": node, "bn": bn_node})
            continue

# print(valid_ConvTranspose_list)
param_dict = {}
for node in valid_ConvTranspose_list:
    conv = node["conv"]
    bn = node["bn"]
    # find params
    param_name = list(conv.input) + list(bn.input)
    for i, initializer in enumerate(all_initializer):
        if initializer.name in param_name:
            param_dict[initializer.name] = onnx.numpy_helper.to_array(initializer)
# print(param_dict)
for node in valid_ConvTranspose_list:
    conv = node["conv"]
    bn = node["bn"]

    bn_eps = bn.attribute[0].f
    bn_mom = bn.attribute[1].f

    bn_w = param_dict[bn.input[1]]  # [Cout, ]
    bn_b = param_dict[bn.input[2]]  # [Cout, ]
    bn_mean = param_dict[bn.input[3]]  # [Cout, ]
    bn_var = param_dict[bn.input[4]]  # [Cout, ]

    conv_w = param_dict[conv.input[1]]  # [Cin, Cout, H, W]
    if len(conv.input) > 2:
        conv_b = param_dict[conv.input[2]]
    else:
        conv_b = np.zeros_like(bn_b)  # [Cout, ]
    conv_w_tran = conv_w.transpose(1, 0, 2, 3)

    Cout = conv_w_tran.shape[0]
    conv_w_reshape = conv_w_tran.reshape([Cout, -1])
    w_bn = np.diag(bn_w / (np.sqrt(bn_eps + bn_var)))
    new_conv_w = np.matmul(w_bn, conv_w_reshape).reshape(conv_w_tran.shape).transpose(1, 0, 2, 3)
    bn_b_tmp = bn_b - (np.multiply(bn_w, bn_mean) / (np.sqrt(bn_eps + bn_var)))
    new_conv_b = np.matmul(bn_w, conv_b) + bn_b_tmp

    new_node = onnx.helper.make_node(
        name=conv.name+'_bn',
        op_type="ConvTranspose",
        inputs=[conv.input[0], conv.name+'_bn.weights', conv.name+'_bn.bias'],
        outputs=[bn.output[0]],
        dilations=conv.attribute[0].ints,
        group=conv.attribute[1].i,
        kernel_shape=conv.attribute[2].ints,
        pads=conv.attribute[3].ints,
        strides=conv.attribute[4].ints
    )
    initializer_w = onnx.helper.make_tensor(
        name=conv.name+'_bn.weights',
        data_type=onnx.helper.TensorProto.DataType.FLOAT,
        dims=new_conv_w.shape,
        vals=new_conv_w.tobytes(),
        raw=True
    )
    initializer_b = onnx.helper.make_tensor(
        name=conv.name+'_bn.bias',
        data_type=onnx.helper.TensorProto.DataType.FLOAT,
        dims=new_conv_b.shape,
        vals=new_conv_b.tobytes(),
        raw=True
    )

    model.graph.initializer.append(initializer_w)
    model.graph.initializer.append(initializer_b)

    # insert node
    for i, node in enumerate(all_node):
        if conv.name == node.name:
            model.graph.node.insert(i, new_node)
            break
    # clean node
    model.graph.node.remove(conv)
    model.graph.node.remove(bn)

onnx.checker.check_model(model)
onnx.save(model, '../onnx/models/backbone_fuse.onnx')

graph = so.graph_from_file('../onnx/models/backbone_fuse.onnx')
graph = so.clean(graph)
so.check(graph)
so.graph_to_file(graph, '../onnx/models/backbone_fuse.onnx')

Signed-off-by: Takeshi Watanabe <take-cheeze@users.noreply.github.com>
@take-cheeze
take-cheeze requested review from a team as code owners February 26, 2026 08:45
@sonarqubecloud

Copy link
Copy Markdown

take-cheeze pushed a commit to take-cheeze/onnx-optimizer that referenced this pull request Jul 23, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants