wehub-resource-sync
94 lines
3.7 KiB
Python
94 lines
3.7 KiB
Python
# The fastai DataLoader is a drop-in replacement for Pytorch's;
|
|
# no code changes are required other than changing the import line
|
|
from fastai.data.load import DataLoader
|
|
import torch
|
|
from torch import nn
|
|
from torch.optim import SGD
|
|
import torch.nn.functional as F
|
|
from torchvision.transforms import Compose, ToTensor, Normalize
|
|
from torchvision.datasets import MNIST
|
|
|
|
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
|
|
from ignite.metrics import Accuracy, Loss
|
|
from ignite.utils import setup_logger
|
|
|
|
class Net(nn.Module):
|
|
def __init__(self):
|
|
super(Net, self).__init__()
|
|
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
|
|
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
|
|
self.conv2_drop = nn.Dropout2d()
|
|
self.fc1 = nn.Linear(320, 50)
|
|
self.fc2 = nn.Linear(50, 10)
|
|
|
|
def forward(self, x):
|
|
x = F.relu(F.max_pool2d(self.conv1(x), 2))
|
|
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
|
|
x = x.view(-1, 320)
|
|
x = F.relu(self.fc1(x))
|
|
x = F.dropout(x, training=self.training)
|
|
x = self.fc2(x)
|
|
return F.log_softmax(x, dim=-1)
|
|
|
|
def get_data_loaders(train_batch_size, val_batch_size):
|
|
data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
|
|
train_loader = DataLoader(
|
|
MNIST(download=True, root=".", transform=data_transform, train=True), batch_size=train_batch_size, shuffle=True)
|
|
val_loader = DataLoader(
|
|
MNIST(download=False, root=".", transform=data_transform, train=False), 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, val_loader = get_data_loaders(train_batch_size, val_batch_size)
|
|
model = Net()
|
|
device = "cpu"
|
|
if torch.cuda.is_available(): device = "cuda"
|
|
|
|
model.to(device) # Move model before creating optimizer
|
|
optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
|
|
criterion = nn.NLLLoss()
|
|
trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
|
|
trainer.logger = setup_logger("trainer")
|
|
|
|
val_metrics = {"accuracy": Accuracy(), "nll": Loss(criterion)}
|
|
evaluator = create_supervised_evaluator(model, metrics=val_metrics, device=device)
|
|
evaluator.logger = setup_logger("evaluator")
|
|
|
|
desc = "ITERATION - loss: {:.2f}"
|
|
pbar = tqdm(initial=0, leave=False, total=len(train_loader), desc=desc.format(0))
|
|
|
|
@trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
|
|
def log_training_loss(engine):
|
|
pbar.desc = desc.format(engine.state.output)
|
|
pbar.update(log_interval)
|
|
|
|
@trainer.on(Events.EPOCH_COMPLETED)
|
|
def log_training_results(engine):
|
|
pbar.refresh()
|
|
evaluator.run(train_loader)
|
|
metrics = evaluator.state.metrics
|
|
avg_accuracy = metrics["accuracy"]
|
|
avg_nll = metrics["nll"]
|
|
tqdm.write(
|
|
"Training Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}".format(
|
|
engine.state.epoch, avg_accuracy, avg_nll))
|
|
|
|
@trainer.on(Events.EPOCH_COMPLETED)
|
|
def log_validation_results(engine):
|
|
evaluator.run(val_loader)
|
|
metrics = evaluator.state.metrics
|
|
avg_accuracy = metrics["accuracy"]
|
|
avg_nll = metrics["nll"]
|
|
tqdm.write(
|
|
"Validation Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}".format(
|
|
engine.state.epoch, avg_accuracy, avg_nll))
|
|
pbar.n = pbar.last_print_n = 0
|
|
|
|
@trainer.on(Events.EPOCH_COMPLETED | Events.COMPLETED)
|
|
def log_time(engine):
|
|
tqdm.write(
|
|
"{} took {} seconds".format(trainer.last_event_name.name, trainer.state.times[trainer.last_event_name.name]))
|
|
|
|
trainer.run(train_loader, max_epochs=epochs)
|
|
|