Training loop with grads

bengio-train-with-grads.py

%matplotlib inline

def plot_loss(trn_loss, val_loss=None, title="Loss Curves"):
    plt.figure(figsize=(10, 6))
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.title(title)

    legends = []

    assert len(trn_loss) % 1000 == 0
    plt.plot(torch.tensor(trn_loss).view(-1, 1000).mean(dim=1))
    legends.append("train loss")

    if val_loss:
        plt.plot(torch.tensor(val_loss).view(-1, 1000).mean(dim=1))
        legends.append("val loss")

    plt.legend(legends)
    plt.show()

# Training loop
trn_loss = []
val_loss = []

def train(device, max_step, X_trn, Y_trn, X_val, Y_val, batch_size, g, model, params, lr, trn_loss, val_loss, with_grad=False):
    grads = {}

    for i in range(max_step):
        ix = torch.randint(0, X_trn.shape[0], (batch_size,), generator=g, device=device)
        x = X_trn[ix]

        # Forward pass
        for layer in model:
            x = layer(x)

        # Compute loss
        loss = F.cross_entropy(x, Y_trn[ix])
        trn_loss.append(loss.item())

        # Retain all gradients for visualization
        if with_grad:
            for layer in model:
                layer.out.retain_grad()

        # Zero gradients to prevent accumulation
        for p in params:
            p.grad = None

        # Backpropagation
        loss.backward()

        if i > 80000:
            lr = 1e-4

        # Update params
        for p in params:
            p.data += -lr * p.grad

        # Copy gradients for visualizations
        if with_grad and i % 20000 == 0:
            for j, layer in enumerate(model):
                if j not in grads:
                    grads[j] = {}
                grads[j][i] = layer.out.grad.cpu().tolist()

        # Validation
        with torch.no_grad():
            ix = torch.randint(0, X_val.shape[0], (batch_size,), generator=g, device=device)
            x = X_val[ix]

            for layer in model:
                x = layer(x)

            loss = F.cross_entropy(x, Y_val[ix])
            val_loss.append(loss.item())

        if i % 10000 == 0:
            print(f"step {i:7d} | train loss {trn_loss[-1]:.4f} | val loss {val_loss[-1]:.4f}")

    return trn_loss, val_loss, grads