GoodWin/Deep-Multi-scale

1#! /usr/bin/env python32# -*- coding: utf-8 -*-3# File   : batchnorm_reimpl.py4# Author : acgtyrant5# Date   : 11/01/20186#7# This file is part of Synchronized-BatchNorm-PyTorch.8# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch9# Distributed under MIT License.10 11import torch12import torch.nn as nn13import torch.nn.init as init14 15__all__ = ['BatchNorm2dReimpl']16 17 18class BatchNorm2dReimpl(nn.Module):19    """20    A re-implementation of batch normalization, used for testing the numerical21    stability.22 23    Author: acgtyrant24    See also:25    https://github.com/vacancy/Synchronized-BatchNorm-PyTorch/issues/1426    """27    def __init__(self, num_features, eps=1e-5, momentum=0.1):28        super().__init__()29 30        self.num_features = num_features31        self.eps = eps32        self.momentum = momentum33        self.weight = nn.Parameter(torch.empty(num_features))34        self.bias = nn.Parameter(torch.empty(num_features))35        self.register_buffer('running_mean', torch.zeros(num_features))36        self.register_buffer('running_var', torch.ones(num_features))37        self.reset_parameters()38 39    def reset_running_stats(self):40        self.running_mean.zero_()41        self.running_var.fill_(1)42 43    def reset_parameters(self):44        self.reset_running_stats()45        init.uniform_(self.weight)46        init.zeros_(self.bias)47 48    def forward(self, input_):49        batchsize, channels, height, width = input_.size()50        numel = batchsize * height * width51        input_ = input_.permute(1, 0, 2, 3).contiguous().view(channels, numel)52        sum_ = input_.sum(1)53        sum_of_square = input_.pow(2).sum(1)54        mean = sum_ / numel55        sumvar = sum_of_square - sum_ * mean56 57        self.running_mean = (58                (1 - self.momentum) * self.running_mean59                + self.momentum * mean.detach()60        )61        unbias_var = sumvar / (numel - 1)62        self.running_var = (63                (1 - self.momentum) * self.running_var64                + self.momentum * unbias_var.detach()65        )66 67        bias_var = sumvar / numel68        inv_std = 1 / (bias_var + self.eps).pow(0.5)69        output = (70                (input_ - mean.unsqueeze(1)) * inv_std.unsqueeze(1) *71                self.weight.unsqueeze(1) + self.bias.unsqueeze(1))72 73        return output.view(channels, batchsize, height, width).permute(1, 0, 2, 3).contiguous()74 75