කණ්ඩායම-නාලිකාසාමාන්යකරණය

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

33class BatchChannelNorm(Module):

43    def __init__(self, channels: int, groups: int,
44                 eps: float = 1e-5, momentum: float = 0.1, estimate: bool = True):

52        super().__init__()

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

63        self.channel_norm = ChannelNorm(channels, groups, eps)

65    def forward(self, x):
66        x = self.batch_norm(x)
67        return self.channel_norm(x)

70class EstimatedBatchNorm(Module):

91    def __init__(self, channels: int,
92                 eps: float = 1e-5, momentum: float = 0.1, affine: bool = True):

99        super().__init__()
100
101        self.eps = eps
102        self.momentum = momentum
103        self.affine = affine
104        self.channels = channels

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

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

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

123        x_shape = x.shape

125        batch_size = x_shape[0]

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

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

134        if self.training:

136            with torch.no_grad():

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

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

145                var = mean_x2 - mean ** 2

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

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

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

167        return x_norm.view(x_shape)

170class ChannelNorm(Module):

177    def __init__(self, channels, groups,
178                 eps: float = 1e-5, affine: bool = True):

185        super().__init__()
186        self.channels = channels
187        self.groups = groups
188        self.eps = eps
189        self.affine = affine

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

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

207        x_shape = x.shape

209        batch_size = x_shape[0]

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

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

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

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

224        var = mean_x2 - mean ** 2

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

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

237        return x_norm.view(x_shape)