How to drop running stats to default value for Norm layer in pyTorch?

user_9140159 0
Viewed 264

I trained model on some images. Now to fit similar dataset but with another colors I want to load this model but also i want to drop all running stats from Batchnorm layers (set them to default value, like totally untrained). What parameters should i reset? Simple model looks like this

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.conv0 = nn.Conv2d(3, 3, 3, padding = 1)
        self.norm = nn.BatchNorm2d(3)
        self.conv = nn.Conv2d(3, 3, 3, padding = 1)
    
def forward(self, x):
        x = self.conv0(x)
        x = self.norm(x)
        return self.conv(x)

net = Net()
##or for pretrained it will be
##net = torch.load('net.pth')

def drop_to_default():
    for m in net.modules():
        if type(m) == nn.BatchNorm2d:
           ####???####

drop_to_default()
pytorch
1 Answers

Simplest way to do that is to run reset_running_stats() method on BatchNorm objects:

def drop_to_default():
    for m in net.modules():
        if type(m) == nn.BatchNorm2d:
           m.reset_running_stats()

Below is this method's source code:

def reset_running_stats(self) -> None:
      if self.track_running_stats:
            # running_mean/running_var/num_batches... are registered at runtime depending
            # if self.track_running_stats is on
          self.running_mean.zero_()  # Zero (neutral) mean
          self.running_var.fill_(1)  # One (neutral) variance
          self.num_batches_tracked.zero_()  # Number of batches tracked

You can see the source code here, _NormBase class.

user_108864200
Related