BatchNorm Folding

Fold BatchNorm layers into the preceding convolutions

BN_Folder merges every BatchNorm2d into the convolution that feeds it: the batch-norm layer becomes an Identity and the convolution weights absorb the affine transform.

from fastai.vision.all import *
from fasterai.misc.all import *

Get the data

path = untar_data(URLs.PETS)
files = get_image_files(path/"images")

def label_func(f): return f[0].isupper()

dls = ImageDataLoaders.from_name_func(path, files, label_func, item_tfms=Resize(64))

Train the model

learn = Learner(dls, resnet18(num_classes=2), metrics=accuracy)
learn.fit_one_cycle(5)

epoch	train_loss	valid_loss	accuracy	time
0	0.618243	0.596936	0.707713	00:03
1	0.568171	0.571559	0.723951	00:03
2	0.510462	0.571453	0.752368	00:03
3	0.444737	0.422476	0.801759	00:03
4	0.374354	0.400150	0.814614	00:03

Fold

fold() requires the model in eval mode: the running statistics it folds in have to be frozen.

learn.model.eval()
folded = BN_Folder().fold(learn.model)

folded.layer1

Sequential(
  (0): BasicBlock(
    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (bn1): Identity()
    (relu): ReLU(inplace=True)
    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (bn2): Identity()
  )
  (1): BasicBlock(
    (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (bn1): Identity()
    (relu): ReLU(inplace=True)
    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (bn2): Identity()
  )
)

bn1 and bn2 are now Identity, and the convolutions carry a bias they did not have before: it holds the batch-norm shift.

What changed

def count_parameters(model): return sum(p.numel() for p in model.parameters())

n_bn = sum(1 for m in learn.model.modules() if isinstance(m, nn.BatchNorm2d))
print(f'{n_bn} BatchNorm2d layers')
print(f'original {count_parameters(learn.model):,}')
print(f'folded   {count_parameters(folded):,}')
print(f'removed  {count_parameters(learn.model) - count_parameters(folded):,}')

20 BatchNorm2d layers
original 11,177,538
folded   11,172,738
removed  4,800

4,800 parameters fewer: the 20 BatchNorm2d layers held 4,800 weights and 4,800 biases, and the convolutions gained the 4,800 biases in exchange.

x, _ = dls.train.one_batch()
learn.model.eval()
with torch.no_grad():
    delta = (learn.model(x) - folded(x)).abs().max().item()

print(f'max |difference| over one training batch of {x.shape[0]}: {delta:.2e}')

max |difference| over one training batch of 64: 4.48e-04

The folded model is an algebraic rewrite of the original, so in exact arithmetic the two are the same function. In float32 the operations are reordered; over one batch of 64 taken from the training loader the outputs differ by at most 4.48e-04.

def wilson(k, n, z=1.96):
    "k out of n and its Wilson 95% interval, as a printable string"
    p, d = k / n, 1 + z**2 / n
    c, h = (p + z**2 / (2 * n)) / d, z * ((p * (1 - p) / n + z**2 / (4 * n**2))**0.5) / d
    return f'{k}/{n} = {100 * p:.2f}% [{100 * (c - h):.2f}, {100 * (c + h):.2f}]'

n = len(dls.valid_ds)
for name, m in [('original', learn.model), ('folded', folded)]:
    lrn = Learner(dls, m, metrics=accuracy)
    with lrn.no_bar(), lrn.no_logging():
        acc = lrn.validate()[1]
    print(f'{name:10s} {wilson(round(acc * n), n)}')

original   1204/1478 = 81.46% [79.40, 83.36]
folded     1204/1478 = 81.46% [79.40, 83.36]

agree = 0
with torch.no_grad():
    for xb, yb in dls.valid:
        agree += (learn.model(xb).argmax(1) == folded(xb).argmax(1)).sum().item()

print(f'{agree}/{n} predictions agree between the two models')

1478/1478 predictions agree between the two models

Equal counts would not by themselves mean equal predictions, so the two models are also compared image by image: they pick the same class on all 1478.

Both models get 1204 of the 1478 validation images right: 81.46%, Wilson 95% [79.40, 83.36]. Single run, one seed; this page measures no latency.

Summary

Call	What it gives you
`BN_Folder().fold(model)`	a copy of the model with every `BatchNorm2d` folded into the convolution before it and replaced by `Identity`

On this ResNet-18 the fold removed the 20 batch-norm layers and 4,800 parameters, moved the outputs by at most 4.48e-04 on one training batch, and left every one of the 1478 validation predictions unchanged. fold() asserts eval mode, so call it after training.

Summary

See Also