Some ideas for BatchNorm Implementation

[gboduljak]

Given that BatchNorm was already requested (#204), I was thinking of how to implement it. Here is my proposal.

We can implement generic BatchNorm class, receiving the number of input axes, reduction axes and num_features. It performs the normalization and keeps track of the running stats.

import mlx.core as mx
from mlx.nn.layers.base import Module

class BatchNorm(Module):
    def __init__(self,
                 num_axes: int,
                 reduction_axes: Tuple[int],
                 num_features: int,
                 eps: float = 1e-5,
                 momentum: float = 0.1,
                 affine: bool = True,
                 track_running_stats: bool = True
                 ):
        r"""Applies batch normalization [1] on the inputs.
        [1]: https://arxiv.org/abs/1502.03167
        """
        super().__init__()
        self.num_axes = num_axes
        self.reduction_axes = reduction_axes
        self.num_features = num_features
        param_shape = tuple(
            [1 if ax in self.reduction_axes else self.num_features for ax in range(
                self.num_axes)]
        )
        if affine:
            self.weight = mx.ones(param_shape)
            self.bias = mx.zeros(param_shape)
        if track_running_stats:
            self.running_mean = mx.zeros(num_features)
            self.running_var = mx.ones(num_features)

        self.track_running_stats = track_running_stats
        self.eps = eps
        self.momentum = momentum

    def __call__(self, x: mx.array):
        if not self._training:
            return (x - self.running_mean) * mx.rsqrt(self.running_var + self.eps)

        batch_mean = mx.mean(x, axis=self.reduction_axes, keepdims=True)
        batch_var = mx.mean(
            (x - batch_mean) ** 2,  # naive
            axis=self.reduction_axes,
            keepdims=True
        )
        x_norm = (x - batch_mean) * mx.rsqrt(batch_var + self.eps)

        if self.track_running_stats:
            self.running_mean = (1.0 - self.momentum) * \
                self.running_mean + self.momentum * batch_mean
            self.running_var = (1.0 - self.momentum) * \
                self.running_var + self.momentum * batch_var

        return self.weight * x_norm + self.bias

Now, we can handle BatchNorm1d and BatchNorm2d as follows:

class BatchNorm1d(BatchNorm):
    def __init__(self,
                 num_features: int,
                 eps: float = 1e-5,
                 momentum: float = 0.1,
                 affine: bool = True,
                 track_running_stats: bool = True):
        super(BatchNorm1d, self).__init__(
            num_axes=2,
            reduction_axes=tuple([0]),
            num_features=num_features,
            eps=eps,
            momentum=momentum,
            affine=affine,
            track_running_stats=track_running_stats
        )

class BatchNorm2d(BatchNorm):
    def __init__(self,
                 num_features: int,
                 eps: float = 1e-5,
                 momentum: float = 0.1,
                 affine: bool = True,
                 track_running_stats: bool = True):
        super(BatchNorm2d, self).__init__(
            num_axes=4,
            reduction_axes=(0, 2, 3),
            num_features=num_features,
            eps=eps,
            momentum=momentum,
            affine=affine,
            track_running_stats=track_running_stats
        )

The following tests are passing,

import mlx.core as mx
import mlx.nn.layers.normalization as norm
import numpy as np
import torch
import torch.nn as tnn

batch_size = 4
num_features = 32
num_iters = 5

our_bn_1d = norm.BatchNorm1d(num_features)
torch_bn_1d = tnn.BatchNorm1d(num_features)

our_bn_1d.train()
torch_bn_1d.train()

with torch.no_grad():
    x_mx = (mx.arange(
        batch_size * num_features) * (1.0)).reshape((batch_size, num_features))
    x_torch = (torch.arange(
        batch_size * num_features) * (1.0)).reshape((batch_size, num_features))

    for _ in range(num_iters):
        print('before')
        print(x_mx)
        print('after')
        x_mx_bn = our_bn_1d(x_mx)
        x_torch_bn = torch_bn_1d(x_torch)
        print(x_mx_bn)
        print(x_torch_bn)
        print('---')
        assert np.allclose(x_mx_bn, x_torch_bn, atol=1e-5)
        x_mx = x_mx_bn * x_mx
        x_torch = x_torch_bn * x_torch

import mlx.core as mx
import mlx.nn.layers.normalization as norm
import numpy as np
import torch
import torch.nn as tnn

batch_size = 4
num_channels = 32
h = 16
w = 16
num_iters = 5

our_bn_2d = norm.BatchNorm2d(num_channels)
torch_bn_2d = tnn.BatchNorm2d(num_channels)

our_bn_2d.train()
torch_bn_2d.train()

with torch.no_grad():
    x_mx = (mx.arange(
        batch_size * num_channels * h * w) * (1.0)).reshape((batch_size, num_channels, h, w))
    x_torch = (torch.arange(
        batch_size * num_channels * h * w) * (1.0)).reshape((batch_size, num_channels, h, w))

    for _ in range(num_iters):
        print('before')
        print(x_mx)
        print('after')
        x_mx_bn = our_bn_2d(x_mx)
        x_torch_bn = torch_bn_2d(x_torch)
        print(x_mx_bn)
        print(x_torch_bn)
        print('---')
        assert np.allclose(x_mx_bn, x_torch_bn, atol=1e-5)
        x_mx = x_mx_bn * x_mx
        x_torch = x_torch_bn * x_torch

Please share some thoughts. If you like the proposal, I can convert this into a PR.

[awni]

I think it's nice! FYI there is already a PR #217 for Batch Norm. Maybe you could review / collaborate on it?

[gboduljak]

Thanks for pointing this out. I think that my discussion post and the PR were concurrent. I reviewed the PR.