Reproducing GroupNorm (Running Mean & Var)

Roy_Eyono1: <p>I am attempting to replicate the exact implementation of GroupNorm, but upon testing, I’m off by approximately <code>1e-2</code> when randomly generating a tensor and feeding it through my implementation of GroupNorm and pytorch’s GroupNorm function.</p>
<p>I suspect that there is a running average implementation in GroupNorm that accounts for the difference? Here is my code:</p>
<p>(PS: Ignore the name of the class, it is not relevant to the implementation )</p>
<pre><code class="lang-auto"># This implementation of GroupNormalization comes from the original paper:
# Figure 3 in https://arxiv.org/pdf/1803.08494.pdf

import torch
import torch.nn as nn

class SubtractiveOnlyGroupNorm(nn.Module):
    def __init__(self, num_groups, num_channels, affine=True, eps=1e-5, momentum=1e-4):
        super().__init__()
        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps
        self.affine = affine
        self.running_mean = 0
        self.running_var = 0
        self.momentum = momentum
        if self.affine:
            self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1))
            self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

    def forward(self, x):
        # Reshape the input tensor so that the spatial dimensions and channels are grouped together
        # We assume that the input has shape (batch_size, num_channels, height, width)
        batch_size, num_channels, height, width = x.size()
        x = x.view(batch_size, self.num_groups, num_channels // self.num_groups, height, width)
        
        # Calculate the mean and variance for each group
        # mean = x.mean(dim=(2, 3, 4), keepdim=True)
        # var = x.var(dim=(2, 3, 4), keepdim=True)
        mean = torch.mean(x, dim=(2,3,4), keepdim=True)
        var = torch.var(x, dim=(2,3,4), unbiased=True, keepdim=True)

        self.running_mean =  self.momentum * self.running_mean + (1 - self.momentum) * mean
        self.running_var = self.momentum * self.running_var + (1 - self.momentum) * var

        # Apply normalization
        x = (x - self.running_mean) / torch.sqrt(self.running_var + self.eps)
        
        # Reshape the normalized tensor back to its original shape
        x = x.view(batch_size, num_channels, height, width)
        
        # Apply the learned weight and bias
        if self.affine:
            x = x * self.weight + self.bias
        
        return x


# main function
if __name__ == "__main__":
    x = torch.randn(32, 3, 28, 28) # batch_size, num_channels, height, width
    norm1 = nn.GroupNorm(1, 3, affine=False, eps=1e-5)
    norm2 = SubtractiveOnlyGroupNorm(1, 3, affine=False, eps=1e-5)
    y1 = norm1(x)[3][0][0][1]
    y2 = norm2(x)[3][0][0][1]
    print(y1, y2)
    print((norm1(x) == norm2(x)).all())
</code></pre>

ptrblck: <p>From the <a href="https://pytorch.org/docs/stable/generated/torch.nn.GroupNorm.html#torch.nn.GroupNorm">docs</a>:</p>
<blockquote>
<p>The standard-deviation is calculated via the biased estimator, equivalent to torch.var(input, unbiased=False).</p>
</blockquote>
<blockquote>
<p>This layer uses statistics computed from input data in both training and evaluation modes.</p>
</blockquote>
<p>After fixing this the error matches the expected precision limit of <code>float32</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">class SubtractiveOnlyGroupNorm(nn.Module):
    def __init__(self, num_groups, num_channels, affine=True, eps=1e-5, momentum=1e-4):
        super().__init__()
        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps
        self.affine = affine
        self.running_mean = 0
        self.running_var = 0
        self.momentum = momentum
        if self.affine:
            self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1))
            self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

    def forward(self, x):
        # Reshape the input tensor so that the spatial dimensions and channels are grouped together
        # We assume that the input has shape (batch_size, num_channels, height, width)
        batch_size, num_channels, height, width = x.size()
        x = x.view(batch_size, self.num_groups, num_channels // self.num_groups, height, width)

        mean = torch.mean(x, dim=(2,3,4), keepdim=True)
        var = torch.var(x, dim=(2,3,4), unbiased=False, keepdim=True)
        # Apply normalization
        x = (x - mean) / torch.sqrt(var + self.eps)
        
        # Reshape the normalized tensor back to its original shape
        x = x.view(batch_size, num_channels, height, width)
        
        # Apply the learned weight and bias
        if self.affine:
            x = x * self.weight + self.bias
        
        return x


# main function
if __name__ == "__main__":
    x = torch.randn(32, 3, 28, 28) # batch_size, num_channels, height, width
    norm1 = nn.GroupNorm(1, 3, affine=False, eps=1e-5)
    norm2 = SubtractiveOnlyGroupNorm(1, 3, affine=False, eps=1e-5)
    y1 = norm1(x)[3][0][0][1]
    y2 = norm2(x)[3][0][0][1]
    print(y1, y2)
    # tensor(1.6263) tensor(1.6263)
    print((norm1(x) - norm2(x)).abs().max())
    # tensor(4.7684e-07)
</code></pre>

Roy_Eyono1: <p>Thank you so much, it sounds like that did the trick!</p>
<p>I have one more question:</p>
<ol>
<li>Can you confirm that there is no running average and var computation in the Pytorch groupnorm implementation?</li>
</ol>

ptrblck: <p>Yes, this is how I understand:</p>
<blockquote>
<p>This layer uses statistics computed from input data in both training and evaluation modes.</p>
</blockquote>

Roy_Eyono1: <p>In case anyone reading this thread is interested in the GroupNorm code, I made a small Github <a href="https://github.com/RoyHEyono/Pytorch-GroupNorm" rel="noopener nofollow ugc">repo</a>.</p>