Cudnn 算子融合

MikeLiu5775
116 阅读3分钟

这里我们以视觉网络中常见的 conv+ bn 说起

一、推理融合

截屏2025-08-16 18.13.46.png

1. BatchNorm 的公式

对于通道 cc 的输入特征图 xx，BatchNorm 在训练时公式为：

y=γ⋅x−μσ2+ϵ+βy = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

其中：

  • μ,σ2\mu, \sigma^2：均值和方差（训练后会存成推理时固定的 running_mean 和 running_var）
  • γ,β\gamma, \beta：可学习的缩放和偏移参数
  • ϵ\epsilon：防止除零的小常数

2. Conv 的公式

假设卷积层的权重为 W，偏置为 b，卷积计算为：

z=W*x+b

其中 * 表示卷积操作。

3. Conv + BN 融合原理

把 BN 的线性变换合并进卷积：

  1. BN 先把卷积结果标准化：
y=γ⋅(W∗x+b)−μσ2+ϵ+βy = \gamma \cdot \frac{(W*x + b) - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta
  1. 重新整理成：
y=(γσ2+ϵ)(W∗x)+(γσ2+ϵ(b−μ)+β)y = \left(\frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}\right) (W*x) + \left(\frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} (b - \mu) + \beta \right)

  1. 定义新的卷积参数：

    • 新权重：
    W′=W⋅γσ2+ϵW' = W \cdot \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}
    • 新偏置：
    b′=γσ2+ϵ(b−μ)+βb' = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} (b - \mu) + \beta

最终得到等效的 Conv-only 层：

y=W′∗x+b′y = W' * x + b'

4. 融合的好处

  • 减少推理开销：少了一层 BN 计算和访存。
  • 简化计算图：利于部署到 TensorRT、TVM、ONNXRuntime、NCNN 等推理框架。

二、训练融合

前向融合

我们继续以该图分析

截屏2025-08-16 18.13.46.png

bn 训练前向

给定的某一通道内：

1、批统计：求均值 标准差

μB=1m∑i=1mxi,σB2=1m∑i=1m(xi−μB)2\mu_B=\frac{1}{m}\sum_{i=1}^m x_i,\qquad \sigma_B^2=\frac{1}{m}\sum_{i=1}^m (x_i-\mu_B)^2

2、求归一化

x^i=xi−μBσB2+ε,yi=γ x^i+β\hat{x}_i=\frac{x_i-\mu_B}{\sqrt{\sigma_B^2+\varepsilon}},\qquad y_i=\gamma\,\hat{x}_i+\beta

以上两步都需要遍历当前通道内NHW个点

卷积前向：

假设卷积层的权重为 W，偏置为 b，卷积计算为：

z=W*x+b

其中 * 表示卷积操作。

这里我们以cuda 编程模型来考虑，为了追求性能，卷积和 bn 训练的滑块无法做到统一。所以对于训练融合我们换个思路，将 bn 训练第一步和第二步在计算图上拆为两个计算节点，归一化节点结合 relu，conv 进行融合。第一部分不进行融合，这样计算图可以粗略概括为，conv+bn+relu+bn==>conv+stats+(scale+bias+relu+conv)

ConvBNfprop.png

下载.png

反向融合

反向融合首先要明白反向传播的原理，其本质就是链式求导

bn 的反向传播也主要分为两部分

第一步求 db ds

第二步求 A B C 三个系数并进行运算得出 dinput

推导如下

首先明白最终目标是求得 ∂L∂xi\displaystyle \frac{\partial L}{\partial x_i}。

回顾前向公式

μ=1m∑j=1mxj,σ2=1m∑j=1m(xj−μ)2,x^i=xi−μσ2+ε,yi=γx^i+β.\mu=\tfrac{1}{m}\sum_{j=1}^m x_j,\quad \sigma^2=\tfrac{1}{m}\sum_{j=1}^m (x_j-\mu)^2,\quad \hat{x}_i=\frac{x_i-\mu}{\sqrt{\sigma^2+\varepsilon}},\quad y_i=\gamma\hat{x}_i+\beta.

所以反向公式分为三部分

∂L∂xi=∂L∂x^i∂x^i∂xi+∂L∂μ∂μ∂xi+∂L∂σ2∂σ2∂xi\frac{\partial L}{\partial x_i} = {\frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial x_i}} + {\frac{\partial L}{\partial \mu}\frac{\partial \mu}{\partial x_i}} + {\frac{\partial L}{\partial \sigma^2}\frac{\partial \sigma^2}{\partial x_i}}
1、∂L∂x^i∂x^i∂xi=∂L∂yi∂yi∂x^i∂x^i∂xi=∂L∂yiγ 1、 \frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial x_i} = \frac{\partial L}{\partial y_i} \frac{\partial y_i}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial x_i} = \frac{\partial L}{\partial y_i}\gamma
2、∂L∂μ∂μ∂xi2、\frac{\partial L}{\partial \mu}\frac{\partial \mu}{\partial x_i}
对于∂L∂μ,由两部分组成，一个是x^i中μ,另一部分是σ2中的μ，所以对于\frac{\partial L}{\partial \mu},由两部分组成，一个是\hat{x}_i 中\mu,另一部分是\sigma^2中的\mu，所以
∂L∂μ=∑i=1m∂L∂x^i∂x^i∂μ+∑i=1m∂L∂x^i∂x^i∂σ2∂σ2∂μ,\frac{\partial L}{\partial \mu}=\sum_{i=1}^m\frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial \mu}+\sum_{i=1}^m\frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial \sigma^2}\frac{\partial \sigma^2}{\partial \mu},

首先我们看下

∂σ2∂μ=−2m∑j=1m(xj−μ)这个导数恒等于0，所以∂L∂μ的值由前半部分∑i=1m∂L∂x^i∂x^i∂μ决定\frac{\partial \sigma^2}{\partial \mu}=\tfrac{-2}{m}\sum_{j=1}^m (x_j-\mu) 这个导数恒等于 0，所以\frac{\partial L}{\partial \mu}的值由前半部分\sum_{i=1}^m\frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial \mu}决定

∑i=1m∂L∂x^i∂x^i∂μ=∑i=1m∂L∂x^i−1σ2+ε\sum_{i=1}^m\frac{\partial L}{\partial \hat{x}_i}\frac{\partial \hat{x}_i}{\partial \mu}=\sum_{i=1}^m\frac{\partial L}{\partial \hat{x}_i}\frac{-1}{\sqrt{\sigma^2+\varepsilon}}
又由于∂μ∂xi=1m，所以∂L∂μ∂μ∂xi=1m∑∂L∂x^i−1σ 又由于 \frac{\partial \mu}{\partial x_i}=\frac{1}{m}，所以\frac{\partial L}{\partial \mu}\frac{\partial \mu}{\partial x_i}=\frac{1}{m}\sum\frac{\partial L}{\partial \hat{x}_i}\frac{-1}{\sigma}
3、继续分析∂L∂σ2∂σ2∂xi3、继续分析{\frac{\partial L}{\partial \sigma^2}\frac{\partial \sigma^2}{\partial x_i}}

把 x^i=(xi−μ)inv_std,把inv_std看作1σ2+ε\hat{x}_i=(x_i-\mu)\text{inv\_std} ,把 {\text {inv\_std}} 看作 \frac{1}{\sqrt{\sigma^2+\varepsilon}} 看作 x^i\hat{x}_i 对 σ2\sigma^2 的函数：

∂x^i∂σ2=(xi−μ)⋅∂inv_std∂σ2=(xi−μ)⋅(−12)(σ2+ε)−3/2.\frac{\partial \hat{x}_i}{\partial \sigma^2} =(x_i-\mu)\cdot\frac{\partial \text{inv\_std}}{\partial \sigma^2} =(x_i-\mu)\cdot\Big(-\tfrac{1}{2}\Big)(\sigma^2+\varepsilon)^{-3/2}.

于是

∂L∂σ2=∑idx^i∂x^i∂σ2=(−12)(σ2+ε)−3/2∑idx^i(xi−μ).\frac{\partial L}{\partial \sigma^2} =\sum_i d\hat{x}_i\frac{\partial \hat{x}_i}{\partial \sigma^2} =\Big(-\tfrac{1}{2}\Big)(\sigma^2+\varepsilon)^{-3/2}\sum_i d\hat{x}_i(x_i-\mu).

又由于

∂σ2∂xi=2m(xi−μ)\frac{\partial \sigma^2}{\partial x_i}=\frac{2}{m}(x_i-\mu)

所以

∂L∂σ2∂σ2∂xi=2m(xi−μ)∗(−12)(σ2+ε)−3/2∑jdx^j(xj−μ){\frac{\partial L}{\partial \sigma^2}\frac{\partial \sigma^2}{\partial x_i}}=\frac{2}{m}(x_i-\mu)*\Big(-\tfrac{1}{2}\Big)(\sigma^2+\varepsilon)^{-3/2}\sum_j d\hat{x}_j(x_j-\mu)

将以上求得的三部分相加

  ∂L∂xi=inv_stdm(m dx^i−S1−x^i S2)  \boxed{\;\frac{\partial L}{\partial x_i} =\frac{\text{inv\_std}}{m}\Big(m\,d\hat{x}_i - S_1 - \hat{x}_i\,S_2\Big)\;}
S1(db)=∑jdx^j,S2(ds)=∑jdx^j x^j.S_1(db)=\sum_j d\hat{x}_j,\qquad S_2(ds)=\sum_j d\hat{x}_j\,\hat{x}_j.

relu 的反向传播 relu 前向公式

y=ReLU(x)=max⁡(0,x)y = \text{ReLU}(x) = \max(0, x)

对输入 xx 求导：

∂y∂x={1如果 x>00如果 x<0\frac{\partial y}{\partial x} = \begin{cases} 1 & \text{如果 } x > 0 \\ 0 & \text{如果 } x < 0 \end{cases}

假设我们有损失函数 LL，它依赖于 ReLU 的输出 yy，我们要求

∂L∂x\frac{\partial L}{\partial x}

由链式法则：

∂L∂x=∂L∂y⋅∂y∂x\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x}

其中：

  • ∂L∂y\frac{\partial L}{\partial y} 来自上游（前一层传下来的梯度）
  • ∂y∂x\frac{\partial y}{\partial x} 是 ReLU 的导数（0 或 1）

所以：

∂L∂x={∂L∂y如果 x>00如果 x≤0\frac{\partial L}{\partial x} = \begin{cases} \frac{\partial L}{\partial y} & \text{如果 } x > 0 \\ 0 & \text{如果 } x \leq 0 \end{cases}

在实际框架里，比如 PyTorch，ReLU backward 常写成：

grad_input = grad_output * (x > 0)

所以整个bn+relu+卷积的反向流程如图所示，这里类似前向分为两个算子，卷积 relu 为一个融合算子，求db ds为另一个算子

image.png

三、融合方式

预编译模式 不灵活但是算子可定制优化

动态生成编译模式 灵活但是算子不好定制优化性能略差

avatar
文章
阅读
粉丝