nn.BatchNorm中affine参数的作用

在PyTorch的nn.BatchNorm2d中，affine参数决定是否在批归一化（Batch Normalization）过程中引入可学习的缩放和平移参数。

BN层的公式如下，

affine参数决定是否在批归一化之后应用一个可学习的线性变换，即缩放和平移。具体来说，如果 affine=True，批归一化层会有两个额外的可学习参数：缩放参数γ（初始值为1）和平移参数β（初始值为0），归一化后的输出变为，

若affine=False，则γ和β都为None。相关的源码定义如下，

self.affine = affine
if self.affine:
    self.weight = Parameter(torch.empty(num_features, **factory_kwargs))
    self.bias = Parameter(torch.empty(num_features, **factory_kwargs))
else:
    self.register_parameter("weight", None)
    self.register_parameter("bias", None)

if self.affine:
    init.ones_(self.weight)
    init.zeros_(self.bias)

以下是简单的代码示例，

import torch
import torch.nn as nn


bn_affine = nn.BatchNorm2d(num_features=3, affine=True)
print("bn_affine γ:", bn_affine.weight)
print("bn_affine β:", bn_affine.bias)
"""
bn_affine.weight
Parameter containing:
tensor([1., 1., 1.], requires_grad=True)

bn_affine.bias
Parameter containing:
tensor([0., 0., 0.], requires_grad=True)
"""


bn_no_affine = nn.BatchNorm2d(num_features=3, affine=False)
print("bn_no_affine γ:", bn_affine.weight)
print("bn_no_affine β:", bn_affine.bias)
"""
None
None
"""