|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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() |
|
|
|
|
