torch.nn.utils.fuse_linear_bn_eval

torch.nn.utils.fuse_linear_bn_eval(linear, bn)

将一个线性模块 (linear module) 和一个 BatchNorm 模块融合成一个新的线性模块。

参数

• linear (torch.nn.Linear) – 一个线性模块。

• bn (torch.nn.modules.batchnorm._BatchNorm) – 一个 BatchNorm 模块。

返回值

融合后的线性模块。

返回类型

torch.nn.Linear

注意

linear 和 bn 都必须处于评估模式 (eval mode)，并且 bn 的运行缓冲区 (running buffers) 必须已计算。