<gh_stars>1-10
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from torch import nn


class FrozenBatchNorm2d(nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters
    are fixed
    """

    def __init__(self, num_features, eps=0):
        # in pytorch batchnorm layer, eps is 1e-5. We assume the default init
        # weight has merged 1e-5 and thus ignored this. Here, we set it as 0 by
        # default to make it compatible with the original maskrcnn
        # implementation, but it is noting that if the init weight comes from
        # pytorch or torch vision, make sure eps is correctly set.
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(num_features))
        self.register_buffer("bias", torch.zeros(num_features))
        self.register_buffer("running_mean", torch.zeros(num_features))
        self.register_buffer("running_var", torch.ones(num_features))
        self.eps = eps
        self.num_features = num_features

    def forward(self, x):
        # Cast all fixed parameters to half() if necessary
        if x.dtype == torch.float16:
            self.weight = self.weight.half()
            self.bias = self.bias.half()
            self.running_mean = self.running_mean.half()
            self.running_var = self.running_var.half()

        scale = self.weight * (self.running_var + self.eps).rsqrt()
        bias = self.bias - self.running_mean * scale
        scale = scale.reshape(1, -1, 1, 1)
        bias = bias.reshape(1, -1, 1, 1)
        return x * scale + bias

    def extra_repr(self):
        return '{}, eps={}'.format(len(self.weight), self.eps)

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import torch
from torch import nn


class FrozenBatchNorm2d(nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters
    are fixed
    """

    def __init__(self, num_features, eps=0):
        # in pytorch batchnorm layer, eps is 1e-5. We assume the default init
        # weight has merged 1e-5 and thus ignored this. Here, we set it as 0 by
        # default to make it compatible with the original maskrcnn
        # implementation, but it is noting that if the init weight comes from
        # pytorch or torch vision, make sure eps is correctly set.
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(num_features))
        self.register_buffer("bias", torch.zeros(num_features))
        self.register_buffer("running_mean", torch.zeros(num_features))
        self.register_buffer("running_var", torch.ones(num_features))
        self.eps = eps
        self.num_features = num_features

    def forward(self, x):
        # Cast all fixed parameters to half() if necessary
        if x.dtype == torch.float16:
            self.weight = self.weight.half()
            self.bias = self.bias.half()
            self.running_mean = self.running_mean.half()
            self.running_var = self.running_var.half()

        scale = self.weight * (self.running_var + self.eps).rsqrt()
        bias = self.bias - self.running_mean * scale
        scale = scale.reshape(1, -1, 1, 1)
        bias = bias.reshape(1, -1, 1, 1)
        return x * scale + bias

    def extra_repr(self):
        return '{}, eps={}'.format(len(self.weight), self.eps)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. BatchNorm2d where the batch statistics and the affine parameters     are fixed # in pytorch batchnorm layer, eps is 1e-5. We assume the default init # weight has merged 1e-5 and thus ignored this. Here, we set it as 0 by # default to make it compatible with the original maskrcnn # implementation, but it is noting that if the init weight comes from # pytorch or torch vision, make sure eps is correctly set. # Cast all fixed parameters to half() if necessary