batchnorm融合

conv-batch_norm 模块可以融合成一个conv模块．
参考链接：https://nenadmarkus.com/p/fusing-batchnorm-and-conv/
测试代码：

# conv 与bn融合ｄｅｍｏ
# https://nenadmarkus.com/p/fusing-batchnorm-and-conv/

import ipdb
import torch
import torchvision

def fuse_conv_and_bn(conv, bn):
    #对于bn,这里共使用５个参数．
    # ２个可训练参数:weight(gamma)和bias(beta)，均值方差(2次)：running_mean(u), running_var(e)和容差eps
    # bias为bn的输入通道（上一层的输出通道）个数．即，针对上层的out_channels每个channel分别求均值方差．
    # init
    print(f'original conv, inc:{conv.in_channels},otc {conv.out_channels}, k:{conv.kernel_size}')
    fusedconv = torch.nn.Conv2d(
        conv.in_channels,
        conv.out_channels,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        bias=True
    )
    #
    # prepare filters, w_conv的行数应该等于上一层的输出通道数
    w_conv = conv.weight.clone().view(conv.out_channels, -1) #conv.in_channels:3, conv.out_channels:64, \
                                                            #conv.weight:(64, 3, 7, 7) --> (64, 147)
    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps+bn.running_var)))# bn.running_var.shape=(64,)
    fusedconv.weight.copy_( torch.mm(w_bn, w_conv).view(fusedconv.weight.size()) )
    #
    # prepare spatial bias
    if conv.bias is not None:
        b_conv = conv.bias
    else:
        b_conv = torch.zeros( conv.weight.size(0) )
    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
    fusedconv.bias.copy_( b_conv + b_bn )
    #
    # we're done
    return fusedconv
    
class DummyModule(nn.Module):
    def __init__(self):
        super(DummyModule, self).__init__()

    def forward(self, x):
        return x
        
def convert():
   # torch.set_grad_enabled(False)
    x = torch.randn(1, 3, 256, 256)
    rn18 = torchvision.models.resnet18(pretrained=True)
    rn18.eval()
    with torch.no_grad():
        ori_out = rn18(x)
        new_bn = fuse_conv_and_bn(rn18.conv1, rn18.bn1)
        rn18.conv1 = new_bn
        rn18.bn1 = DummyModule()
        new_out = rn18(x)
        max_diff = (ori_out - new_out).max()
        print(f'max diff:{max_diff.item()}')

if __name__ == '__main__':
    convert()

参考资料：https://github.com/qinjian623/pytorch_toys
https://github.com/Ironteen/Batch-Normalization-fusion