pytorch融合conv和bn

为提高推理速度，一般会将卷积层和对应的bn层融合。不是废话，直接上代码。

# 融合前的卷积和bn
# bias 一般为False
conv1 = nn.Conv2d(in_chn, out_chn, 3, padding=1, bias=False)
bn = nn.BatchNorm2d(out_chn)

# 融合函数, 返回融合后的卷积核权重和偏置
def merge_bn(conv, bn):
    kernel = conv.weight.data
    running_mean = bn.running_mean
    running_var = bn.running_var
    gamma = bn.weight.data
    beta = bn.bias.data
    eps = bn.eps
    std = torch.sqrt(running_var+eps)
    t = (gamma/std).view(-1,1,1,1).expand(kernel.size())
    kernel = kernel*t
    bias = beta - running_mean*gamma/std
    return kernel, bias