Running stats with gradient checkpointing

[vovaf709]

According to patch_batchnorm source code if layer collecting running stats (e.g. BatchNorm) is checkpointed it will accumulate statistics only when grad is enabled (on backward pass). This induces inconsistency:

torch.manual_seed(1337)
seq = nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4))
torch.manual_seed(1337)
seq_checkpointed = checkpoint_wrapper(nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4)))

inp = torch.randn(2, 4, 16, 16)

seq(inp)
seq_checkpointed(inp)

seq[1].running_mean == seq_checkpointed[1].running_mean
tensor([False, False, False, False])

I think this behaviour should be modified to accumulate statistics at 1-st forward pass or at least mentioned in docs

[min-xu-ai]

Thanks for reporting. I slightly modified your code to demonstrate how it works:

import torch
from torch import nn
from fairscale.nn.checkpoint import checkpoint_wrapper

torch.manual_seed(1337)
seq = nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4))
torch.manual_seed(1337)
seq_checkpointed = checkpoint_wrapper(nn.Sequential(nn.Conv2d(4, 4, 3), nn.BatchNorm2d(4)))

inp = torch.randn(2, 4, 16, 16).requires_grad_(True)

out = seq(inp)
out_ck = seq_checkpointed(inp)
torch.testing.assert_close(out, out_ck)

out.sum().backward()
out_ck.sum().backward()


print(seq[1].running_mean)
print(seq_checkpointed[1].running_mean)
torch.testing.assert_close(seq[1].running_mean, seq_checkpointed[1].running_mean)
#tensor([False, False, False, False])

As you can see, you need to run the backward pass to make the running_mean match. Just forward is not enough. Checkpoint_wrapper is used for training. Only doing the forward pass does not make sense IMHO. With backward pass, the stats is matching correctly.

[vovaf709]

There is one kaggle trick - you can run multiple forward passes on test set to adapt running stats to it. This is weird case but still)

[min-xu-ai]

Oh I see. That’s interesting! Do you have example code or pseudo code for it?

[vovaf709]

Code for this trick? If yes then it is as simple as

model.train()
for (X, y), _ in zip(test_loader, range(n_iter)):
    model(X)

[min-xu-ai]

I see. Then after this loop you proceed with normal training for 1 epoch or the whole training N epochs?

[vovaf709]

No, I run this loop on the test set (on which I want to get the highest target metric in competition) after the whole training. The idea is to adapt BN statistics to the test set which can have slightly different distribution

[vovaf709]

I think I can come up with a solution on the next week, ok?

[nyngwang]

@vovaf709 ok