ML Architecture Comparison

def train(model, train_loader, val_loader, criterion,
          optimizer, epochs, device, save_path,
          epoch_offset=0, patience=10):

    best_val_loss = float("inf")
    patience_counter = 0

    for epoch in range(epochs):
        train_loss, train_correct, train_total = 0.0, 0, 0
        model.train()

        for inputs, labels in tqdm(train_loader):
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()

        val_loss, val_acc = evaluate(model, val_loader, criterion, device)

        wandb.log({
            "train/loss": train_loss / len(train_loader),
            "train/acc": train_correct / train_total,
            "val/loss": val_loss,
            "val/acc": val_acc,
            "epoch": epoch + 1 + epoch_offset
        })

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            torch.save(model.state_dict(), save_path)
        else:
            patience_counter += 1
            if patience_counter >= patience:
                break

    return epoch + 1