Convolution和Batch normalization的融合

理論推算

當前CNN卷積層的基本組成單元標配：Conv + BN +ReLU 三個子模塊。但其實在網絡的推理階段，可以將BN層的運算融合到Conv層中，减少運算量，加速推理。本質上是修改了卷積核的參數，在不增加Conv層計算量的同時，略去了BN層的計算量。公式推導如下。

conv層的參數

BN層的參數

假設輸入為x，則x->Conv->BN的輸出便是：

做個簡單的公式變形：

代碼實現

在實際使用時，首先要定比特conv和bn的比特置，根據實際情况進行替換或者删除BN層。在本次實施例中，以開源分割模型庫https://github.com/qubvel/segmentation_models.pytorch為案例進行融合實驗，對BN層進行了替換。

class Conv2dReLU(nn.Sequential):
    def __init__(
            self,
            in_channels,
            out_channels,
            kernel_size,
            padding=0,
            stride=1,
            use_batchnorm=True,
    ):

        if use_batchnorm == "inplace" and InPlaceABN is None:
            raise RuntimeError(
                "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
                + "To install see: https://github.com/mapillary/inplace_abn"
            )

        conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            bias=not (use_batchnorm),
        )
        relu = nn.ReLU(inplace=True)

        if use_batchnorm == "inplace":
            bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
            relu = nn.Identity()

        elif use_batchnorm and use_batchnorm != "inplace":
            bn = nn.BatchNorm2d(out_channels)

        else:
            bn = nn.Identity()

        super(Conv2dReLU, self).__init__(conv, bn, relu)

from turtle import forward
from torch.fx.experimental.optimization import fuse
import torch
import torch.nn as nn
import time

import segmentation_models_pytorch.base.modules as md

from utils.torchUtils import fuse_conv_and_bn


def fuseModel(model):  # fuse model Conv2d() + BatchNorm2d() layers
    for m in model.modules():
        if isinstance(m, (md.Conv2dReLU)) and isinstance(m[1], (nn.BatchNorm2d)):
            m[0] = fuse_conv_and_bn(m[0], m[1])  # update conv
            m[1] = nn.Identity()
            count += 1
    return model