torch.nn.utils.fuse_linear_bn_weights

torch.nn.utils.fuse_linear_bn_weights(linear_w, linear_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b)

将线性模块参数和 BatchNorm 模块参数融合到新的线性模块参数中。

参数

• linear_w (torch.Tensor) – 线性权重。

• linear_b (Optional[torch.Tensor]) – 线性偏置。

• bn_rm (torch.Tensor) – BatchNorm 运行均值。

• bn_rv (torch.Tensor) – BatchNorm 运行方差。

• bn_eps (float) – BatchNorm epsilon。

• bn_w (torch.Tensor) – BatchNorm 权重。

• bn_b (torch.Tensor) – BatchNorm 偏置。

返回

融合后的线性权重和偏置。

返回类型

Tuple[torch.nn.Parameter, torch.nn.Parameter]