BatchNormalization.py 文件源码

def __init__(self, nOutput, eps=1e-5, momentum=0.1, affine=True):
        super(BatchNormalization, self).__init__()
        assert nOutput != 0

        self.affine = affine
        self.eps = eps
        self.train = True
        self.momentum = momentum
        self.running_mean = torch.zeros(nOutput)
        self.running_var = torch.ones(nOutput)

        self.save_mean = None
        self.save_std = None

        if self.affine:
           self.weight = torch.Tensor(nOutput)
           self.bias = torch.Tensor(nOutput)
           self.gradWeight = torch.Tensor(nOutput)
           self.gradBias = torch.Tensor(nOutput)
           self.reset()
        else:
           self.weight = None
           self.bias = None
           self.gradWeight = None
           self.gradBias = None