Skip to content

frozen_batchnorm

frozen_batchnorm

Frozen BatchNorm2d implementation taken from DETR.

FrozenBatchNorm2d 

FrozenBatchNorm2d(dimension)

Bases: Module

BatchNorm2d where the batch statistics and the affine parameters are fixed.

Initialize with dimension equal to the channel dimension.

Source code in src/versatil/models/layers/frozen_batchnorm.py 
11
12
13
14
15
16
17
def __init__(self, dimension: int):
    """Initialize with dimension equal to the channel dimension."""
    super().__init__()
    self.register_buffer("weight", torch.ones(dimension))
    self.register_buffer("bias", torch.zeros(dimension))
    self.register_buffer("running_mean", torch.zeros(dimension))
    self.register_buffer("running_var", torch.ones(dimension))

forward 

forward(x)

Forward pass for input tensor with shape (B, C, H, W).

Source code in src/versatil/models/layers/frozen_batchnorm.py 
43
44
45
46
47
48
49
50
51
52
def forward(self, x: torch.Tensor):
    """Forward pass for input tensor with shape (B, C, H, W)."""
    w = self.weight.reshape(1, -1, 1, 1)  # Shape: (1, C, 1, 1)
    b = self.bias.reshape(1, -1, 1, 1)
    rv = self.running_var.reshape(1, -1, 1, 1)
    rm = self.running_mean.reshape(1, -1, 1, 1)
    eps = 1e-5
    scale = w * (rv + eps).rsqrt()
    bias = b - rm * scale
    return x * scale + bias