train.py

import torch


def train(model, trainloader, criterion, optimizer, epoch, device):
    model.train()
    train_loss = 0
    correct = 0
    total = 0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()

        if batch_idx % 100 == 0:
            print(f'\rEpoch {epoch}, Batch {batch_idx}, Loss: {train_loss/(batch_idx+1):.4f}, Accuracy: {100.*correct/total:.2f}%', end='', flush=True)


def test(model, testloader, criterion, epoch, best_acc, checkpoint_path, device):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(testloader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    print(f'\rEpoch {epoch}, Test Loss: {test_loss/(batch_idx+1):.4f}, Test Accuracy: {100.*correct/total:.2f}%')
    if 100.*correct/total > best_acc:
        best_acc = 100.*correct/total
        if checkpoint_path is not None:
            torch.save(model.state_dict(), checkpoint_path)
            print(f'Saved best model with acc {100.*correct/total}')
    return best_acc