Quick start¶
Link to quick start jupyter notebook.
Simple example¶
Define a PyTorch model.
import torch
from torch import nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self, n_classes, p_dropout=0.5):
super().__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d(p=p_dropout)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, n_classes)
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 x
2. Define a argus.model.Model
with nn_module
, optimizer
, loss
attributes. Each value must be a class
or function that returns object (torch.nn.Module
for loss and nn_module, torch.optim.Optimizer
for optimizer).
import argus
class MnistModel(argus.Model):
nn_module = Net
optimizer = torch.optim.SGD
loss = torch.nn.CrossEntropyLoss
3. Create an instance of MnistModel
with the specified parameters. Net will be initialized like
Net(n_classes=10, p_dropout=0.1)
. The same logic is applied for the optimizer torch.optim.SGD(lr=0.01)
.
Loss will be created without any arguments torch.nn.CrossEntropyLoss()
. The model will use the CPU.
params = {
'nn_module': {'n_classes': 10, 'p_dropout': 0.1},
'optimizer': {'lr': 0.01},
'device': 'cpu'
}
model = MnistModel(params)
Download MNIST dataset. Create validation and training PyTorch data loaders.
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Normalize
from torchvision.datasets import MNIST
data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
train_mnist_dataset = MNIST(download=True, root="mnist_data",
transform=data_transform, train=True)
val_mnist_dataset = MNIST(download=False, root="mnist_data",
transform=data_transform, train=False)
train_loader = DataLoader(train_mnist_dataset,
batch_size=64, shuffle=True)
val_loader = DataLoader(val_mnist_dataset,
batch_size=128, shuffle=False)
5. Define some callbacks and start training the model for 50 epochs.
As metrics, you can use instances of classes inherit from argus.metrics.Metric
or they names argus.metrics.Metric.name
.
from argus.callbacks import MonitorCheckpoint, EarlyStopping, ReduceLROnPlateau
callbacks = [
MonitorCheckpoint(dir_path='mnist', monitor='val_accuracy', max_saves=3),
EarlyStopping(monitor='val_accuracy', patience=9),
ReduceLROnPlateau(monitor='val_accuracy', factor=0.5, patience=3)
]
model.fit(train_loader,
val_loader=val_loader,
num_epochs=50,
metrics=['accuracy'], # or [argus.metrics.CategoricalAccuracy()]
callbacks=callbacks)
Load the model from the best checkpoint.
from pathlib import Path
model_path = Path("mnist/").glob("*.pth")
model_path = sorted(model_path)[-1]
print(f"Load model: {model_path}")
loaded_model = argus.load_model(model_path)
print(loaded_model)
More flexibility¶
Argus can help you simplify the experiments with different architectures, losses, and optimizers. Let’s define a
argus.model.Model
with two models via a dictionary. If you want to use PyTorch losses and optimizers, it’s not
necessary to define them in the argus model.
from torchvision.models import resnet18
class FlexModel(argus.Model):
nn_module = {
'net': Net,
'resnet18': resnet18
}
Create a model instance. Parameters for nn_module is a tuple where the first element is a name, second is init arguments. PyTorch losses and optimizers can be selected by a string with a class name.
params = {
'nn_module': ('resnet18', {
'pretrained': False,
'num_classes': 1
}),
'optimizer': ('Adam', {'lr': 0.01}),
'loss': 'CrossEntropyLoss',
'device': 'cuda'
}
model = FlexModel(params)
Argus allows managing different parts combinations of a pipeline.
class MoreFlexModel(argus.Model):
nn_module = {
'net': Net,
'resnet18': resnet18
}
optimizer = {
'SGD': torch.optim.SGD,
'adam_w': torch.optim.AdamW
}
loss = {
'BCE': nn.BCEWithLogitsLoss,
'cross_entropy': nn.CrossEntropyLoss,
'nll': nn.NLLLoss
}
prediction_transform = {
'sigmoid': nn.Sigmoid,
'Softmax': nn.Softmax
}
params = {
'nn_module': ('resnet18', {
'pretrained': False,
'num_classes': 1
}),
'optimizer': ('adam_w', {
'lr': 0.01,
'weight_decay': 0.042
}),
'loss': ('BCE', {'reduction': 'sum'}),
'prediction_transform': 'sigmoid',
'device': 'cuda'
}
model = MoreFlexModel(params)
See also
If you need more flexibility you can:
Override methods of
argus.model.Model
. For example,argus.model.Model.train_step()
andargus.model.Model.val_step()
. See Train and val steps guide for details.Create a custom
argus.callbacks.Callback
.Use a custom
argus.metrics.Metric
.