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
| 16 | |
| 17 | |
| 18 | class BatchNorm2dReimpl(nn.Module): |
| 19 | """ |
| 20 | A re-implementation of batch normalization, used for testing the numerical |
| 21 | stability. |
| 22 | |
| 23 | Author: acgtyrant |
| 24 | See also: |
| 25 | https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/14 |
| 26 | """ |
| 27 | |
| 28 | def __init__(self, num_features, eps=1e-5, momentum=0.1): |
| 29 | super().__init__() |
| 30 | |
| 31 | self.num_features = num_features |
| 32 | self.eps = eps |
| 33 | self.momentum = momentum |
| 34 | self.weight = nn.Parameter(torch.empty(num_features)) |
| 35 | self.bias = nn.Parameter(torch.empty(num_features)) |
| 36 | self.register_buffer("running_mean", torch.zeros(num_features)) |
| 37 | self.register_buffer("running_var", torch.ones(num_features)) |
| 38 | self.reset_parameters() |
| 39 | |
| 40 | def reset_running_stats(self): |
| 41 | self.running_mean.zero_() |
| 42 | self.running_var.fill_(1) |
| 43 | |
| 44 | def reset_parameters(self): |
| 45 | self.reset_running_stats() |
| 46 | init.uniform_(self.weight) |
| 47 | init.zeros_(self.bias) |
| 48 | |
| 49 | def forward(self, input_): |
| 50 | batchsize, channels, height, width = input_.size() |
| 51 | numel = batchsize * height * width |
| 52 | input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel) |
| 53 | sum_ = input_.sum(1) |
| 54 | sum_of_square = input_.pow(2).sum(1) |
| 55 | mean = sum_ / numel |
| 56 | sumvar = sum_of_square - sum_ * mean |
| 57 | |
| 58 | self.running_mean = ( |
| 59 | 1 - self.momentum |
| 60 | ) * self.running_mean + self.momentum * mean.detach() |
| 61 | unbias_var = sumvar / (numel - 1) |
| 62 | self.running_var = ( |
| 63 | 1 - self.momentum |
| 64 | ) * self.running_var + self.momentum * unbias_var.detach() |
| 65 | |
| 66 | bias_var = sumvar / numel |
| 67 | inv_std = 1 / (bias_var + self.eps).pow(0.5) |
| 68 | output = (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze( |
| 69 | 1 |
| 70 | ) * self.weight.unsqueeze(1) + self.bias.unsqueeze(1) |
| 71 | |
| 72 | return ( |
| 73 | output.view(channels, batchsize, height, width) |
| 74 | .permute(1, 0, 2, 3) |
| 75 | .contiguous() |
nothing calls this directly
no outgoing calls
no test coverage detected