colossalai.engine

class colossalai.engine.Engine(model, optimizer, criterion=None, gradient_handlers=None, clip_grad_norm=0.0, ophook_list=None, verbose=True)

Basic engine class for training and evaluation. It runs a specific process method step() which is based on the given schedule over each batch of a dataset. It controls a iteration in training.

Parameters

  • model (torch.nn.Module) – The neural network model

  • optimizer (torch.optim.Optimizer) – Optimizer for updating the parameters

  • criterion (torch.nn.modules.loss._Loss, optional) – Loss function for calculating loss

  • gradient_handlers (a list of BaseGradientHandler, optional) – A list of gradient handler used in backward

  • clip_grad_norm (float, optional) – The norm of gradient clipping

  • ophook_list (list) – List of ophook

  • verbose (bool) – whether to display log info

property model

Model attached to the engine

property optimizer

Optimizer attached to the engine

property criterion

Criterion attached to the engine

zero_grad()

Set the gradient of parameters to zero

step()

Execute parameter update

backward(loss)

Start backward propagation given the loss value computed by a loss function

Parameters

loss (torch.Tensor) – Loss value computed by a loss function

backward_by_grad(tensor, grad)

Start backward propagation given the gradient of the output tensor

Parameters

  • tensor (torch.Tensor) – Output tensor

  • grad (torch.Tensor) – Gradient passed back to the output

calc_loss(*args, **kwargs)

Compute the loss value

Parameters

  • args – Args used in criterion function

  • kwargs – Kwargs used in criterion function

Returns

The loss value

Return type

torch.Tensor

train()

Sets the model to training mode.

eval()

Sets the model to evaluation mode.