<jupyter_start><jupyter_text>1. define your dataloader<jupyter_code>transform_list = transforms.Compose([transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.131], std=[0.289])])
dataset = datasets.MNIST(root='/tmp/', train=True, download=True, transform=transform_list)
dataset = Subset(dataset, range(2000))

loader = torch.utils.data.DataLoader(
      dataset=dataset,
      batch_size=500,
      shuffle=False)<jupyter_output><empty_output><jupyter_text>2. define your model<jupyter_code>n_hidden = 2
hidden_size = 10
device = 'cuda'
layers = [nn.Flatten(), nn.Linear(28 * 28, hidden_size), nn.ReLU()] + \
         [nn.Linear(hidden_size, hidden_size), nn.ReLU()] * (n_hidden - 1) + \
         [nn.Linear(hidden_size, 10), nn.LogSoftmax(dim=1)]
model = nn.Sequential(*layers).to(device)<jupyter_output><empty_output><jupyter_text>EWC penaltyWe now compute the matrices coefficients using the generator above. Using different representations we get different performances, and also we use more or less memory. But low memory footprint comes at the price of a less accurate approximation of the FIM, and so the gradient that we get from the penalty can be drastically different.<jupyter_code>from nngeometry.object import PMatKFAC, PMatDiag, PMatBlockDiag, PMatDense, PVector
from nngeometry.metrics import FIM<jupyter_output><empty_output><jupyter_text>Suppose we now train our model on task 1. We want to store the current state of the network using: 1. the current parameter values `v1` 2. the current Fisher Information Matrix `F_XXX`<jupyter_code>v1 = PVector.from_model(model).clone().detach()

F_kfac = FIM(model=model,
             loader=loader,
             representation=PMatKFAC,
             n_output=10,
             variant='classif_logits',
             device='cuda')

F_blockdiag = FIM(model=model,
                  loader=loader,
                  representation=PMatBlockDiag,
                  n_output=10,
                  variant='classif_logits',
                  device='cuda')

F_dense = FIM(model=model,
             loader=loader,
             representation=PMatDense,
             n_output=10,
             variant='classif_logits',
             device='cuda')


F_diag = FIM(model=model,
             loader=loader,
             representation=PMatDiag,
             n_output=10,
             variant='classif_logits',
             device='cuda')<jupyter_output><empty_output><jupyter_text>We can now continue training, using an additional regularizer term, that uses the FIM and the difference betweem the current parameter value and `v1`<jupyter_code>v_current = PVector.from_model(model)<jupyter_output><empty_output><jupyter_text>you can now compute a regularizer scalar object and backward through it<jupyter_code>regularizer_dense = F_dense.vTMv(v_current - v1)
regularizer_dense.backward()
[(p.size(), p.grad) for p in model.parameters()]<jupyter_output><empty_output><jupyter_text>note that in this case all gradients stay at `0`, because `v_current = v1`, so for the purpose of this example, we will modify the value of `v_current`. You can think of this modification as resulting from SGD updates on the first iterations of training on task 2.<jupyter_code>for p in model.parameters():
    p.data.add_(torch.randn_like(p))<jupyter_output><empty_output><jupyter_text>We now recompute the gradients from the regularizer:<jupyter_code>regularizer_dense = F_dense.vTMv(v_current - v1)
regularizer_dense.backward()
[(p.size(), p.grad) for p in model.parameters()]<jupyter_output><empty_output><jupyter_text>Comparison of regularization obtained using different representations<jupyter_code>model.zero_grad()
regularizer_dense = F_dense.vTMv(v_current - v1)
regularizer_dense.backward()
g_dense = torch.cat([p.grad.view(-1) for p in model.parameters()])
model.zero_grad()
regularizer_bd = F_blockdiag.vTMv(v_current - v1)
regularizer_bd.backward()
g_bd = torch.cat([p.grad.view(-1) for p in model.parameters()])
model.zero_grad()
regularizer_kfac = F_kfac.vTMv(v_current - v1)
regularizer_kfac.backward()
g_kfac = torch.cat([p.grad.view(-1) for p in model.parameters()])
model.zero_grad()
regularizer_diag = F_diag.vTMv(v_current - v1)
regularizer_diag.backward()
g_diag = torch.cat([p.grad.view(-1) for p in model.parameters()])
plt.plot(g_kfac[:100].cpu().numpy(), label='kfac')
plt.plot(g_diag[:100].cpu().numpy(), label='diag')
plt.plot(g_dense[:100].cpu().numpy(), label='dense')
plt.plot(g_bd[:100].cpu().numpy(), label='block diagonal')
plt.legend()
plt.title('Compare 100 elements of the gradient (first layer)')
plt.plot(g_kfac[-100:].cpu().numpy(), label='kfac')
plt.plot(g_diag[-100:].cpu().numpy(), label='diag')
plt.plot(g_dense[-100:].cpu().numpy(), label='dense')
plt.plot(g_bd[-100:].cpu().numpy(), label='block diagonal')
# plt.ylim(-.01, .01)
plt.legend()
plt.title('Compare 100 elements of the gradient (last layer)')<jupyter_output><empty_output><jupyter_text>KFAC and Batch norm layers In the case of the KFAC representation, it is not clear what to do with batch norm parameters. If you want to be able to use batch norm layers, you should instead mix representations. You can for instance use the block diagonal representation for batch norm layers, and still use KFAC for the linear layers.Here is how, let us start by defining a new network with batch norm layers:<jupyter_code>layers = [nn.Flatten(), nn.Linear(28 * 28, hidden_size), nn.ReLU()] + \
         [nn.Linear(hidden_size, hidden_size), nn.BatchNorm1d(hidden_size), nn.ReLU()] * (n_hidden - 1) + \
         [nn.Linear(hidden_size, 10), nn.LogSoftmax(dim=1)]
model = nn.Sequential(*layers).to(device)<jupyter_output><empty_output><jupyter_text>We now create 2 different layer collections, one that will contain all linear layers, the other ones only batch norm layers.<jupyter_code>from nngeometry.layercollection import LayerCollection

layer_collection_bn = LayerCollection()
layer_collection_linear = LayerCollection()

for mod in layers:
    if type(mod) in [torch.nn.modules.linear.Linear]:
        layer_collection_linear.add_layer_from_model(model, mod)
    elif type(mod) in [torch.nn.modules.batchnorm.BatchNorm1d]:
        layer_collection_bn.add_layer_from_model(model, mod)<jupyter_output><empty_output><jupyter_text>We now create our 2 representations.Note that both representations do not use the same layer collection, even if they come from the same model.<jupyter_code>F_linear_kfac = FIM(layer_collection=layer_collection_linear,
                    model=model,
                    loader=loader,
                    representation=PMatKFAC,
                    n_output=10,
                    variant='classif_logits',
                    device='cuda')

F_bn_blockdiag = FIM(layer_collection=layer_collection_bn,
                     model=model,
                     loader=loader,
                     representation=PMatBlockDiag,
                     n_output=10,
                     variant='classif_logits',
                     device='cuda')

F_bn_blockdiag.size(), F_linear_kfac.size()
v_current = PVector.from_model(model)
F_linear_kfac.vTMv(v_current), F_bn_blockdiag.vTMv(v_current)<jupyter_output><empty_output><jupyter_text>We can create a new regularizer as the sum of the contributions of linear layers and batch norm layers:<jupyter_code>regularizer = F_linear_kfac.vTMv(v_current) + F_bn_blockdiag.vTMv(v_current)

model.zero_grad()
regularizer.backward()
g_with_bn = torch.cat([p.grad.view(-1) for p in model.parameters()])
plt.plot(g_with_bn.cpu().numpy())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11251646903820817, 'fraction_numerical': 0.01396574440052701, 'mean_word_length': 4.181569965870307, 'pattern_counts': {'":': 0, '<': 53, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '365225', 'license': 'MIT', 'n_tokens_mistral': 2511, 'n_tokens_neox': 2387, 'n_words': 573, 'path': 'examples/FIM for EWC.ipynb', 'repo_name': 'shreyasumbetla/nngeometry', 'size': '7590'}