| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.init as init |
|
|
| __all__ = ["BatchNormReimpl"] |
|
|
|
|
| class BatchNorm2dReimpl(nn.Module): |
| """ |
| A re-implementation of batch normalization, used for testing the numerical |
| stability. |
| |
| Author: acgtyrant |
| See also: |
| https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14 |
| """ |
|
|
| def __init__(self, num_features, eps=1e-5, momentum=0.1): |
| super().__init__() |
|
|
| self.num_features = num_features |
| self.eps = eps |
| self.momentum = momentum |
| self.weight = nn.Parameter(torch.empty(num_features)) |
| self.bias = nn.Parameter(torch.empty(num_features)) |
| self.register_buffer("running_mean", torch.zeros(num_features)) |
| self.register_buffer("running_var", torch.ones(num_features)) |
| self.reset_parameters() |
|
|
| def reset_running_stats(self): |
| self.running_mean.zero_() |
| self.running_var.fill_(1) |
|
|
| def reset_parameters(self): |
| self.reset_running_stats() |
| init.uniform_(self.weight) |
| init.zeros_(self.bias) |
|
|
| def forward(self, input_): |
| batchsize, channels, height, width = input_.size() |
| numel = batchsize * height * width |
| input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel) |
| sum_ = input_.sum(1) |
| sum_of_square = input_.pow(2).sum(1) |
| mean = sum_ / numel |
| sumvar = sum_of_square - sum_ * mean |
|
|
| self.running_mean = ( |
| 1 - self.momentum |
| ) * self.running_mean + self.momentum * mean.detach() |
| unbias_var = sumvar / (numel - 1) |
| self.running_var = ( |
| 1 - self.momentum |
| ) * self.running_var + self.momentum * unbias_var.detach() |
|
|
| bias_var = sumvar / numel |
| inv_std = 1 / (bias_var + self.eps).pow(0.5) |
| output = (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze( |
| 1 |
| ) * self.weight.unsqueeze(1) + self.bias.unsqueeze(1) |
|
|
| return ( |
| output.view(channels, batchsize, height, width) |
| .permute(1, 0, 2, 3) |
| .contiguous() |
| ) |
|
|
