# 批处理信道标准化

25import torch
26from torch import nn
27
28from labml_helpers.module import Module
29from labml_nn.normalization.batch_norm import BatchNorm

## 批量信道规范化

32class BatchChannelNorm(Module):
• channels 是输入中的要素数
• groups 是要素被划分到的组的数量
• eps用于数值稳定性
• momentum 是取指数移动平均线的动量
• estimate 是否使用运行均值和方差作为批次范数
42    def __init__(self, channels: int, groups: int,
43                 eps: float = 1e-5, momentum: float = 0.1, estimate: bool = True):
51        super().__init__()

54        if estimate:
55            self.batch_norm = EstimatedBatchNorm(channels,
56                                                 eps=eps, momentum=momentum)
57        else:
58            self.batch_norm = BatchNorm(channels,
59                                        eps=eps, momentum=momentum)

62        self.channel_norm = ChannelNorm(channels, groups, eps)
64    def forward(self, x):
65        x = self.batch_norm(x)
66        return self.channel_norm(x)

## 预计批次规范化

69class EstimatedBatchNorm(Module):
• channels 是输入中的要素数
• eps用于数值稳定性
• momentum 是取指数移动平均线的动量
• estimate 是否使用运行均值和方差作为批次范数
90    def __init__(self, channels: int,
91                 eps: float = 1e-5, momentum: float = 0.1, affine: bool = True):
98        super().__init__()
99
100        self.eps = eps
101        self.momentum = momentum
102        self.affine = affine
103        self.channels = channels

106        if self.affine:
107            self.scale = nn.Parameter(torch.ones(channels))
108            self.shift = nn.Parameter(torch.zeros(channels))

111        self.register_buffer('exp_mean', torch.zeros(channels))
112        self.register_buffer('exp_var', torch.ones(channels))

x 是形状张量[batch_size, channels, *]* 表示任意数量（可能为 0）的维度。例如，在图像（2D）卷积中，这将是[batch_size, channels, height, width]

114    def forward(self, x: torch.Tensor):

122        x_shape = x.shape

124        batch_size = x_shape[0]

127        assert self.channels == x.shape[1]

130        x = x.view(batch_size, self.channels, -1)

133        if self.training:

135            with torch.no_grad():

138                mean = x.mean(dim=[0, 2])

141                mean_x2 = (x ** 2).mean(dim=[0, 2])

144                var = mean_x2 - mean ** 2

152                self.exp_mean = (1 - self.momentum) * self.exp_mean + self.momentum * mean
153                self.exp_var = (1 - self.momentum) * self.exp_var + self.momentum * var

157        x_norm = (x - self.exp_mean.view(1, -1, 1)) / torch.sqrt(self.exp_var + self.eps).view(1, -1, 1)

162        if self.affine:
163            x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)

166        return x_norm.view(x_shape)

## 频道规范化

169class ChannelNorm(Module):
• groups 是要素被划分到的组的数量
• channels 是输入中的要素数
• eps用于数值稳定性
• affine 是否缩放和移动归一化值
176    def __init__(self, channels, groups,
177                 eps: float = 1e-5, affine: bool = True):
184        super().__init__()
185        self.channels = channels
186        self.groups = groups
187        self.eps = eps
188        self.affine = affine

193        if self.affine:
194            self.scale = nn.Parameter(torch.ones(groups))
195            self.shift = nn.Parameter(torch.zeros(groups))

x 是形状张量[batch_size, channels, *]* 表示任意数量（可能为 0）的维度。例如，在图像（2D）卷积中，这将是[batch_size, channels, height, width]

197    def forward(self, x: torch.Tensor):

206        x_shape = x.shape

208        batch_size = x_shape[0]

210        assert self.channels == x.shape[1]

213        x = x.view(batch_size, self.groups, -1)

217        mean = x.mean(dim=[-1], keepdim=True)

220        mean_x2 = (x ** 2).mean(dim=[-1], keepdim=True)

223        var = mean_x2 - mean ** 2

228        x_norm = (x - mean) / torch.sqrt(var + self.eps)

232        if self.affine:
233            x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)

236        return x_norm.view(x_shape)