Ignite training with online metrics profiling

how_to_run.md

Just run as

python -m cProfile -o run.prof script.py

script.py

from __future__ import print_function
from argparse import ArgumentParser

from torch import nn
from torch.optim import SGD
from torch.utils.data import DataLoader
import torch
import torch.nn.functional as F
from torchvision.transforms import Compose, ToTensor, Normalize
from torchvision.datasets import FakeData
from torchvision.models.resnet import resnet34

from ignite.engine import Events, Engine, _prepare_batch
from ignite.metrics import CategoricalAccuracy, Loss, TopKCategoricalAccuracy


def get_data_loaders(train_batch_size, val_batch_size):
    data_transform = Compose([ToTensor()])

    train_loader = DataLoader(FakeData(size=1000, image_size=(3, 224, 224), num_classes=10, 
                                       transform=data_transform, target_transform=lambda t: t.long()),
                              batch_size=train_batch_size, shuffle=True)

    val_loader = DataLoader(FakeData(size=1000, image_size=(3, 224, 224), num_classes=10, 
                                     transform=data_transform, target_transform=lambda t: t.long()),
                            batch_size=val_batch_size, shuffle=False)
    return train_loader, val_loader


def run(train_batch_size, val_batch_size, epochs, lr, momentum, log_interval):
    train_loader, _ = get_data_loaders(train_batch_size, val_batch_size)
    model = resnet34(pretrained=False, num_classes=10)
    device = 'cpu'

    if torch.cuda.is_available():
        device = 'cuda'

    loss_fn = F.cross_entropy
    optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
    
    if device:
        model.to(device)

    def process_function(engine, batch):
        model.train()
        optimizer.zero_grad()
        x, y = _prepare_batch(batch, device=device)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        loss.backward()
        optimizer.step()
        return [loss.item(), y_pred, y]

    trainer = Engine(process_function)
    
    output_transform = lambda x: [x[1], x[2]]
    metrics={
        'accuracy': CategoricalAccuracy(output_transform=output_transform),
        'topk': TopKCategoricalAccuracy(output_transform=output_transform),
        'xentropy': Loss(loss_fn, output_transform=output_transform)
    }
    for name, metric in metrics.items():
        metric.attach(trainer, name)


    @trainer.on(Events.ITERATION_COMPLETED)
    def log_training_loss(engine):
        iter = (engine.state.iteration - 1) % len(train_loader) + 1
        if iter % log_interval == 0:
            print("Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
                  "".format(engine.state.epoch, iter, len(train_loader), engine.state.output[0]))

    trainer.run(train_loader, max_epochs=epochs)
    print(trainer.state.metrics)
    
run(96, 64, 20, 0.001, 0.1, 10)