PyTorch中的nn.BatchNorm2d

Posted 2021-02-28 dyclown

class _NormBase(Module):   #源码
    """Common base of _InstanceNorm and _BatchNorm"""
    _version = 2
    __constants__ = [‘track_running_stats‘, ‘momentum‘, ‘eps‘,
                     ‘num_features‘, ‘affine‘]

    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
                 track_running_stats=True):
        super(_NormBase, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        if self.affine:
            self.weight = Parameter(torch.Tensor(num_features))
            self.bias = Parameter(torch.Tensor(num_features))
        else:
            self.register_parameter(‘weight‘, None)
            self.register_parameter(‘bias‘, None)
        if self.track_running_stats:
            self.register_buffer(‘running_mean‘, torch.zeros(num_features))
            self.register_buffer(‘running_var‘, torch.ones(num_features))
            self.register_buffer(‘num_batches_tracked‘, torch.tensor(0, dtype=torch.long))
        else:
            self.register_parameter(‘running_mean‘, None)
            self.register_parameter(‘running_var‘, None)
            self.register_parameter(‘num_batches_tracked‘, None)
        self.reset_parameters()

torch.nn.BatchNorm1d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

详细见https://blog.csdn.net/LoseInVain/article/details/86476010