multi_loss_optim

class aitoolbox.torchtrain.multi_loss_optim.MultiLoss(loss_dict, loss_optimizer_map=None, retain_graph_until_last=True)[source]

Bases: object

Multiple loss wrapper for TrainLoop based training

Parameters
  • loss_dict (dict) – dict of loss objects which are used to calculate losses in the TrainLoop

  • loss_optimizer_map (dict or None) – dict mapping the loss name to the corresponding optimizer’s index in the MultiOptimizer. If this parameter is left to None the mapping is automatically created by assigning values from range(len(loss_dict)) as corresponding optimizer indices.

  • retain_graph_until_last (bool) – when calling backward should retain_graph option be enabled for all but last loss tensor

backward(optimizer_idx, iteration, amp_grad_scaler)[source]

Executes backward() for the specific loss based on provided optimizer_idx

Parameters
  • optimizer_idx (int) – index of the current optimizer. Mostly useful when using multiple optimizers. When only a single optimizer is used this parameter can be ignored.

  • iteration (int) – Current iteration index. Not used in the most simple setup but provided in case of more elaborate loss backward logic is devised.

  • amp_grad_scaler (torch.cuda.amp.GradScaler) – AMP GradScaler. If scaler enabled parameter is set to False the loss is still passed to it but it gets returned unscaled so the behaviour is as it is in the case of non-AMP training.

Returns

None

item()[source]
class aitoolbox.torchtrain.multi_loss_optim.MultiOptimizer(optimizer_list)[source]

Bases: object

Multiple optimizer wrapper for TrainLoop based training

Parameters

optimizer_list (list) – list of optimizer objects which are used in the TrainLoop

step(optimizer_idx, iteration, amp_grad_scaler)[source]

Execute step for optimizer at the specified index

Parameters
  • optimizer_idx (int) – index of the current optimizer. Mostly useful when using multiple optimizers. When only a single optimizer is used this parameter can be ignored.

  • iteration (int) – Current iteration index. Not used in the most simple setup but provided in case of more elaborate loss backward logic is devised.

  • amp_grad_scaler (torch.cuda.amp.GradScaler) – AMP GradScaler. If scaler enabled parameter is set to False the optimizer have it’s normal step() method called withoud applying the AMP mandated unscaling beforehand. In this respect the behaviour will be the same as in the non-AMP training.

Returns

None

zero_grad(optimizer_idx, iteration)[source]

Execute zero_grad for optimizer at the specified index

Parameters
  • optimizer_idx (int) – index of the current optimizer. Mostly useful when using multiple optimizers. When only a single optimizer is used this parameter can be ignored.

  • iteration (int) – Current iteration index. Not used in the most simple setup but provided in case of more elaborate loss backward logic is devised.

Returns

None

state_dict()[source]
load_state_dict(state_dict_list)[source]