<jupyter_start><jupyter_text>Data loading<jupyter_code>dataset = FimacDataset(project_dir/'data/interim/renders.hdf5')

split = 0.8
test_size = int(len(dataset) * (1 - split))
train_data, val_data = random_split(dataset, (len(dataset) - test_size, test_size))

train_dataloader = DataLoader(train_data, batch_size=16, shuffle=True)
val_dataloader = DataLoader(val_data, batch_size=32)
help(train_dataloader)
len(train_dataloader.dataset)<jupyter_output><empty_output><jupyter_text>Network definition<jupyter_code>net = TestNet().to(device)
net
eval('TestNet')<jupyter_output><empty_output><jupyter_text>Training loop<jupyter_code>def int_MAE(y, y_hat):
    err = y.type(torch.IntTensor) - y_hat.type(torch.IntTensor)    
    err = err.abs()
    err = err.type(torch.FloatTensor).mean()

    return err.item()
epochs = 5

lr = 0.001
optimizer = torch.optim.SGD(net.parameters(), lr=lr)
criterion = nn.L1Loss()
scaler = GradScaler()

wandb.init(
    project="part-counting",
    entity="brunompac",
    config={
        "learning_rate": lr,
        "epochs": epochs,
        "batch_size": train_dataloader.batch_size,
        "model": type(net).__name__,
        "optmizer": type(optimizer).__name__,
        "loss_func": type(criterion).__name__,
})

wandb.watch(net)

for epoch in range(epochs):
    # train pass
    net.train()
    train_loss = 0
    with torch.set_grad_enabled(True):
        for i, (X, y) in enumerate(train_dataloader):
            X = X.to(device)
            y = y.to(device)

            optimizer.zero_grad()

            with autocast():
                y_hat = net(X)
                loss = criterion(y_hat.squeeze(), y.type(torch.float32))

            scaler.scale(loss).backward()

            train_loss += loss.item() / len(y)  # scales to data size

            scaler.step(optimizer)
            scaler.update()

    wandb.log({
        "train_loss": train_loss,
    }, step=epoch)

    # validation pass
    net.eval()
    val_loss = 0
    val_MAE = 0
    with torch.set_grad_enabled(False):
        for i, (X, y) in enumerate(train_dataloader):
            X = X.to(device)
            y = y.to(device)

            with autocast():
                y_hat = net(X)
                loss = criterion(y_hat.squeeze(), y.type(torch.float32))

            val_loss += loss.item() / len(y)  # scales to data size
            val_MAE += int_MAE(y, y_hat.squeeze()) / len(y)

    wandb.log({
        "val_loss": val_loss,
        "val_MAE": val_MAE,
    }, step=epoch, commit=True)<jupyter_output>[34m[1mwandb[0m: Currently logged in as: [33mbrunompac[0m (use `wandb login --relogin` to force relogin)
{'fraction_non_alphanumeric': 0.13634637295860236, 'fraction_numerical': 0.010634257500949488, 'mean_word_length': 2.86784140969163, 'pattern_counts': {'":': 9, '<': 12, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '440939', 'license': 'CNRI-Python', 'n_tokens_mistral': 992, 'n_tokens_neox': 960, 'n_words': 167, 'path': 'notebooks/02.02-bmp-training-routine-test.ipynb', 'repo_name': 'brunompacheco/part-counting', 'size': '2633'}