Model¶

class argus.model.Model(params: dict)[source]¶

Argus model is an abstraction of a trainer/predictor that uses:

nn_module¶

PyTorch neural network as torch.nn.Module.

Type: torch.nn.Module

optimizer¶

Optimizer as torch.optim.Optimizer.

Type: torch.optim.Optimizer

loss¶

Loss function as torch.nn.Module.

Type: torch.nn.Module

device¶

device as torch.torch.device. The attribute stores a device location of output in the case of multi-GPU mode. To get all devices use argus.model.Model.get_device() method.

Type: torch.device

prediction_transform¶

postprocessing function of predictions as Callable function or object.

Type: Callable

Parameters: params (dict) – A model parameters.

Examples

One can use several ways to initialize argus.model.Model:

Set parameters for each part of the model directly:

class MnistModel(argus.Model):
    nn_module = Net  # torch.nn.Module
    optimizer = torch.optim.SGD
    loss = torch.nn.CrossEntropyLoss

params = {
    'nn_module': {'n_classes': 10, 'p_dropout': 0.1},
    'optimizer': {'lr': 0.01},
    'device': 'cpu'
}

model = MnistModel(params)

2. Set components of the model from multiple options using two elements tuples (component name, the component init arguments):

from torchvision.models import resnet18

class FlexModel(argus.Model):
    nn_module = {
        'net': Net,
        'resnet18': resnet18
    }

params = {
    'nn_module': ('resnet18', {
        'pretrained': False,
        'num_classes': 1
    }),
    'optimizer': ('Adam', {'lr': 0.01}),
    'loss': 'CrossEntropyLoss',
    'device': 'cuda'
}

model = FlexModel(params)

eval()[source]¶: Set the nn_module into eval mode.

fit(train_loader: Iterable, val_loader: Optional[Iterable] = None, num_epochs: int = 1, metrics: Optional[List[Union[argus.metrics.metric.Metric, str]]] = None, metrics_on_train: bool = False, callbacks: Optional[List[argus.callbacks.callback.Callback]] = None, val_callbacks: Optional[List[argus.callbacks.callback.Callback]] = None)[source]¶

Train the argus model.

The method attaches metrics and callbacks to the train and validation process, and performs training itself.

Parameters

train_loader (Iterable) – The train data loader.
val_loader (Iterable, optional) – The validation data loader. Defaults to None.
num_epochs (int, optional) – Number of training epochs to run. Defaults to 1.
metrics (list of argus.metrics.Metric or str, optional) – List of metrics to evaluate. By default, the metrics are evaluated on the validation data (if any) only. Defaults to None.
metrics_on_train (bool, optional) – Evaluate the metrics on train data as well. Defaults to False.
callbacks (list of argus.callbacks.Callback, optional) – List of callbacks to be attached to the training process. Defaults to None.
val_callbacks (list of argus.callbacks.Callback, optional) – List of callbacks to be attached to the validation process. Defaults to None.

get_device() → Union[torch.device, List[torch.device]]¶

Get device or list of devices in case of multi-GPU mode.

Returns: A device or list of devices.
Return type: torch.device or list of torch.device

get_lr() → Union[float, List[float]][source]¶

Get the learning rate from the optimizer.

It could be a single value or a list of values in the case of multiple parameter groups.

Returns: The learning rate value or a list of individual parameter groups learning rate values.
Return type: float or a list of float

get_nn_module() → torch.nn.Module¶

Get nn_module without torch.nn.DataParallel or torch.nn.parallel.DistributedDataParallel.

Returns: nn_module without DP and DDP.
Return type: torch.nn.Module

predict(input)[source]¶

Make a prediction with the given input.

The prediction process consists of the input tensor transferring to the model device, forward pass of the nn_module in eval mode and application of the prediction_transform to the raw prediction output.

Parameters: input (torch.Tensor) – The input tensor to predict with. It will be transferred to the model device. The user is responsible for ensuring that the input tensor shape and type match the model.
Returns: Predictions as the result of the prediction_transform application.
Return type: torch.Tensor or other type

save(file_path: Union[pathlib.Path, str], optimizer_state: bool = False)[source]¶

Save the argus model into a file.

The argus model is saved as a dict:

{
    'model_name': Name of the argus model,
    'params': Argus model parameters dict,
    'nn_state_dict': torch nn_module.state_dict(),
    'optimizer_state_dict': torch optimizer.state_dict()
}

The state_dict is always transferred to CPU before saving.

Parameters

file_path (str or pathlib.Path) – Path to the argus model file.
optimizer_state (bool) – Save optimizer state. Defaults to False.

set_device(device: Union[str, torch.device, List[Union[str, torch.device]]])¶

Move nn_module and loss to the specified device.

If a list of devices is passed, torch.nn.DataParallel will be used. Batch tensors will be scattered on dim 0. The first device in the list is the location of the output. By default, device “cuda” is the GPU training on torch.cuda.current_device().

Example

model.set_device("cuda")
model.set_device(torch.device("cuda"))

model.set_device("cuda:0")
model.set_device(["cuda:2", "cuda:3"])  # Use DataParallel

model.set_device([torch.device("cuda:2"),
                  torch.device("cuda", index=3)])

Parameters: device (str, torch.device or list of devices) – A device or list of devices.

set_lr(lr: Union[float, List[float]])[source]¶

Set the learning rate for the optimizer.

The method allows setting individual learning rates for the optimizer parameter groups as well as setting even learning rate for all parameters.

Parameters

lr (number or list/tuple of numbers) – The learning rate to set. If a single number is provided, all parameter groups learning rates are set to the same value. In order to set individual learning rates for each parameter group, a list or tuple of values with the corresponding length should be provided.

Raises

ValueError – If lr is a list or tuple and its length is not equal to the number of parameter groups.
ValueError – If lr type is not a list, tuple, or number.
AttributeError – If the model is not train_ready (i.e. not all attributes are set).

train(mode: bool = True)[source]¶

Set the nn_module into train mode.

Parameters: mode (bool) – Set train mode, otherwise eval mode. Defaults to True.

train_step(batch, state: argus.engine.engine.State) → dict[source]¶

Perform a single train step.

The method is used by argus.engine.Engine. The train step includes input and target tensor transferring to the model device, forward pass, loss evaluation, backward pass, and the train batch prediction treating with a prediction_transform.

Parameters

torch.Tensors (batch (tuple of 2) – (input, target)): The input and target tensors to process.
state (argus.engine.State) – The argus model state.

Returns

The train step results:

{
    'prediction': The train batch predictions,
    'target': The train batch target data on the model device,
    'loss': The loss function value
}

Return type

dict

val_step(batch, state: argus.engine.engine.State) → dict[source]¶

Perform a single validation step.

The method is used by argus.engine.Engine. The validation step includes input and target tensor transferring to the model device, forward pass, loss evaluation, and the val batch prediction treating with a prediction_transform.

Gradients calculation and the model weights update are omitted, which is the main difference with the train_step() method.

Parameters

torch.Tensors (batch (tuple of 2) – (input, target)): The input data and target tensors to process.
state (argus.engine.State) – The argus model state.

Returns

Default val step results:

{
    'prediction': The val batch predictions,
    'target': The val batch target data on the model device,
    'loss': The loss function value
}

Return type

dict

validate(val_loader: Iterable, metrics: Optional[List[Union[argus.metrics.metric.Metric, str]]] = None, callbacks: Optional[List[argus.callbacks.callback.Callback]] = None) → Dict[str, float][source]¶

Perform a validation.

Parameters

val_loader (Iterable) – The validation data loader.
metrics (list of argus.metrics.Metric or str, optional) – List of metrics to evaluate with the data. Defaults to None.
callbacks (list of argus.callbacks.Callback, optional) – List of callbacks to be attached to the validation process. Defaults to None.

Returns

The metrics dictionary.

Return type

dict

Load argus model¶

argus.model.load_model(file_path: Union[pathlib.Path, str], nn_module: Union[argus.utils.Default, dict, Tuple[str, dict]] = default, optimizer: Union[argus.utils.Default, None, dict, Tuple[str, dict]] = default, loss: Union[argus.utils.Default, None, dict, Tuple[str, dict]] = default, prediction_transform: Union[argus.utils.Default, None, dict, Tuple[str, dict]] = default, device: Union[argus.utils.Default, str, torch.device, List[Union[str, torch.device]]] = default, state_load_func: Callable = <function default_state_load_func>, change_params_func: Callable = Identity(), change_state_dict_func: Callable = <function default_change_state_dict_func>, model_name: Union[argus.utils.Default, str] = default, **kwargs)[source]¶

Load an argus model from a file.

The function allows loading an argus model, saved with argus.model.Model.save(). The model is always loaded in eval mode.

Parameters

file_path (str or pathlib.Path) – Path to the file to load.
nn_module (dict, tuple or str, optional) – Params of the nn_module to replace params in the state.
optimizer (None, dict, tuple or str, optional) – Params of the optimizer to replace params in the state. Optimizer is not created in the loaded model if it is set to None.
loss (None, dict, tuple or str, optional) – Params of the loss to replace params in the state. Loss is not created in the loaded model if it is set to None.
prediction_transform (None, dict, tuple or str, optional) – Params of the prediction_transform to replace params in the state. prediction_transform is not created in the loaded model if it is set to None.
device (str, torch.device or list of devices, optional) – The model device.
state_load_func (function, optional) – Function for loading state from file path.
change_params_func (function, optional) – Function for modification of the loaded params. It takes params from the loaded state as an input and outputs params to use during the model creation.
change_state_dict_func (function, optional) – Function for modification of nn_module and optimizer state dict. Takes nn_state_dict and optimizer_state_dict as inputs and outputs state dicts for the model creation.
model_name (str, optional) – Class name of argus.model.Model. By default uses the name from the loaded state.

Returns

Loaded argus model.

Return type

argus.model.Model

Example

model = ArgusModel(params)
model.save(model_path, optimizer_state=True)

# restarting python...

# ArgusModel class must be already in the scope
model = argus.load_model(model_path, device="cuda:0")

You can find more options how to use load_model here.

Raises

ImportError – If the model is not available in the scope. Often it means that it is not imported or defined.
FileNotFoundError – If the file is not found by the file_path.

API reference

Callbacks