Just run as
python -m cProfile -o run.prof script.py
Just run as
python -m cProfile -o run.prof 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) |