Method forward

imperative/python/megengine/module/batchnorm.py:165–214

(self, inp)

            )

    def forward(self, inp):
        self._check_input_ndim(inp)

        inp_shape = inp.shape
        _ndims = len(inp_shape)
        if _ndims != 4:
            new_shape = Tensor([1, 1, 1, 1], device=inp.device)
            origin_shape = inp_shape
            if _ndims == 2:
                new_shape[:2] = origin_shape[:2]
            elif _ndims == 3:
                new_shape[:3] = origin_shape[:3]
            else:
                raise ValueError(
                    "expected 2D, 3D or 4D input (got {}D input)".format(len(inp_shape))
                )

            inp = inp.reshape(new_shape)

        if self.training and self.track_running_stats:
            exponential_average_factor = self.momentum
        else:
            exponential_average_factor = 0.0  # useless

        _weight = self.weight
        _bias = self.bias

        if self.freeze:
            if _weight is not None:
                _weight = _weight.detach()
            if _bias is not None:
                _bias = _bias.detach()

        output = sync_batch_norm(
            inp,
            self.running_mean,
            self.running_var,
            _weight,
            _bias,
            training=(self.training and not self.freeze)
            or ((self.running_mean is None) and (self.running_var is None)),
            momentum=exponential_average_factor,
            eps=self.eps,
            group=self.group,
        )

        if _ndims != 4:
            output = output.reshape(origin_shape)

        return output


class BatchNorm1d(_BatchNorm):