argus.model

class argus.model.Model(params: dict)[source]
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, and runs the training process.

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, 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_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 floats)

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[str, pathlib.Path])[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()
}

The state_dict is always transferred to cpu prior to saving.

Parameters

file_path (str) – Path to the argus model file.

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 list, tuple, or number.

  • AttributeError – If the model is not train_ready (i.e. not all attributes are set).

train()[source]

Set the nn_module into train mode.

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 transition to the model device, forward pass, loss evaluation, backward pass, and the train batch prediction preparation with a prediction_transform.

Parameters
  • (tuple of 2 torch.Tensors (batch) – (input, target)): The input data 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': 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 transition to the model device, forward pass, loss evaluation, and the train batch prediction preparation with a prediction_transform.

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

Parameters
  • (tuple of 2 torch.Tensors (batch) – (input, target)): The input data 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': Loss function value
}

Return type

dict

validate(val_loader: Optional[Iterable], metrics: Optional[List[argus.metrics.metric.Metric]] = 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, 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[str, pathlib.Path], nn_module=default, optimizer=default, loss=default, prediction_transform=default, device=default, change_params_func=Identity(), change_state_dict_func=Identity(), model_name=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) – Path to the file to load.

  • device (str or torch.device, optional) – Device for the model. Defaults to None.

  • nn_module (dict, tuple or str, optional) – Params of the nn_module to replace params in the state.

  • optimizer (dict, tuple or str, optional) – Params of the optimizer to replace params in the state. Set to None if don’t want to create optimizer in the loaded model.

  • loss (dict, tuple or str, optional) – Params of the loss to replace params in the state. Set to None if don’t want to create loss in the loaded model.

  • prediction_transform (dict, tuple or str, optional) – Params of the prediction_transform to replace params in the state. Set to None if don’t want to create prediction_transform in the loaded model.

  • change_params_func (function, optional) – Function for modification of state params. Takes as input params from the loaded state, outputs params to model creation.

  • change_state_dict_func (function, optional) – Function for modification of nn_module state dict. Takes as input state dict from the loaded state, outputs state dict to model creation.

  • model_name (str) – Class name of argus.model.Model. By default uses name from loaded state.

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.

Returns

Loaded argus model.

Return type

argus.model.Model