Analysing an SGD NN Optimisation Trajectory

Simulate the XOR gate using a neural network (NN). Collect the trajectory of SGD when training the NN and analyse it.

Setup NN

seed = 100

import torch
import torch.nn as nn
import numpy as np
import seaborn as sns
import pandas as pd
import torch.nn.init as init

class XOR(nn.Module):
    def __init__(self):
        super(XOR, self).__init__()

        self.layers = nn.Sequential(
            nn.Linear(2, 2),
            nn.Tanh(),
            nn.Linear(2, 1),
        )
    def forward(self, x):
        return self.layers(x)

def weight_init(m):
    '''
    Usage:
        model = Model()
        model.apply(weight_init)
    '''
    
    if isinstance(m, nn.Linear):
        init.xavier_normal_(m.weight.data)
        init.uniform_(m.bias.data)

def get_weights_biases(model: XOR):
    return [
        model.layers[0].weight[0][0].item(),
        model.layers[0].weight[0][1].item(),
        model.layers[0].weight[1][0].item(),
        model.layers[0].weight[1][1].item(),
        model.layers[0].bias[0].item(),
        model.layers[0].bias[1].item(),
        model.layers[2].weight[0][0].item(),
        model.layers[2].weight[0][1].item(),
        model.layers[2].bias[0].item(),
    ]

device = torch.device("cpu")

x_train = torch.tensor([[0,0], [0,1], [1,1], [1,0]], device=device).float()
y_train = torch.tensor([[0], [1], [0], [1]], device=device).float()


x_val = torch.clone(x_train)
y_val = torch.clone(y_train)

torch.manual_seed(seed)
np.random.seed(seed)

model = XOR()
model.apply(weight_init)
model = model.to(device)

lr = 0.1
n_epochs = 1000

optimizer = torch.optim.SGD(model.parameters(), lr=lr)
loss_fn = nn.BCEWithLogitsLoss()

weights = []
losses = []
for epoch in range(n_epochs):
    model.train()

    logits = model(x_train)

    loss = loss_fn(logits, y_train)

    losses.append(loss.item())
    
    loss.backward()

    optimizer.step()

    weights.append(get_weights_biases(model))
    
    optimizer.zero_grad()
    if (epoch % 50 == 0):
        print("Epoch: {0}, Loss: {1}, ".format(epoch, loss.to("cpu").detach().numpy()))

X = pd.DataFrame(weights)
X.columns = ['w0', 'w1', 'w2', 'w3', 'b0', 'b1', 'w4', 'w5', 'b2']
F = pd.DataFrame()
F['train_loss'] = losses
F['epoch'] = list(F.index)

sample = {
    'name': 'xor-sgd',
    'X': X,
    'F': F.drop(columns='epoch'),
    'N': 'X-index' # use index in X to infer neighbourhood
}

ax = sns.lineplot(F, x='epoch', y='train_loss')

Setup pyXla sample

from pyxla import load_data
import pyxla

load_data(sample)

from matplotlib import rc_file
rc_file('../../../../src/.matplotlibrc')

feat, plot = pyxla.distr_f(sample, title=False)

plot.savefig('distr-f-xor-nn-trained.pdf', bbox_inches='tight')

from matplotlib import rcParams
rcParams['figure.figsize'] = [2.2, 2.2]

corr, imp, plot = pyxla.X_imp(sample, estimator='ridge', alpha=1, n_repeats=30, seed=seed, suptitle=False)

plot.savefig('x-imp-xor-nn-trained.pdf')

from matplotlib import rc_file
rc_file('../../../../src/.matplotlibrc')

corr, plot = pyxla.fdc(sample)
plot.savefig('fdc-xor-nn-trained.pdf', bbox_inches='tight')

corr, plot = pyxla.rdc(sample)

corr, plot = pyxla.pdc(sample)

pyxla.disp_best(sample)

corr, plot = pyxla.nfc(sample)

corr, plot = pyxla.nrc(sample)