Simplified RunningBatchNorm in Lesson 11

At the beginning of Lesson 11, at 2:18 Jeremy shows a simplified (finalized?) RunningBatchNorm class in an updated version of the 07_batchnorm.ipynb notebook. See the screenshot below. I couldn’t find this notebook in the fastai repo.

How can we access the updated version of the 07_batchnorm.ipynb notebook, and why isn’t it in the fastai/course-v3 repo?

2 Likes

Hello,
I stumbled upon this thread while searching for the same thing.

Here is the simplified version from the lecture, that I slightly modified to generalize it to n-dimensions (and also for the corner case bs=1). But at the moment i get much worse results with it, I’m not sure why yet.

class RunningBatchNorm(nn.Module):
    def __init__(self, nf, mom=0.1, eps=1e-5, dims = (0,2,3)):
        super().__init__()
        sp_dims = (1,)*(len(dims)-1)
        self.mom,self.eps,self.dims = mom,eps,dims
        self.mults = nn.Parameter(torch.ones((nf,)+sp_dims))
        self.adds = nn.Parameter(torch.zeros((nf,)+sp_dims))
        self.register_buffer('sums', torch.zeros((1,nf,)+sp_dims))
        self.register_buffer('sqrs', torch.zeros((1,nf,)+sp_dims))
        self.register_buffer('count', tensor(0.))
        self.register_buffer('factor', tensor(0.))
        self.register_buffer('offset', tensor(0.))
        self.batch = 0

    def update_stats(self, x):
        bs,nc,*_ = x.shape
        self.sums.detach_()
        self.sqrs.detach_()
        dims = self.dims
        s = x.sum(dims, keepdim=True)
        ss = (x*x).sum(dims, keepdim=True)
        c = s.new_tensor(x.numel()/nc)
        mom1 = s.new_tensor(1 - (1-self.mom)/math.sqrt(max(bs-1, 1)))
        self.sums.lerp_(s, mom1)
        self.sqrs.lerp_(ss, mom1)
        self.count.lerp_(c, mom1)
        self.batch += bs
        means = self.sums/self.count
        varns = (self.sqrs/self.count).sub_(means*means)
        if bool(self.batch < 20): varns.clamp_min_(0.01)
        self.factor = self.mults / (varns+self.eps).sqrt()
        self.offset = self.adds - means*self.factor

    def forward(self, x):
        if self.training: self.update_stats(x)
        return x*self.factor + self.offset

RunningBatchNorm2d = RunningBatchNorm

class RunningBatchNorm3d(RunningBatchNorm):
    def __init__(self, *args, dims = (0,2,3,4), **kwargs):
        super().__init__(*args, dims=dims, **kwargs)

I also tried this code with fp16, and it’s complaining that lerp_cuda isn’t implemented for half. This is strange because while googling I found those api from Nvidia from cuda in C/C++. Maybe it isn’t supported in pytorch yet?

Has anyone tried lerp with fp16 yet?