colossalai.engine

class colossalai.engine.Engine(model, optimizer, criterion, gradient_handlers=None, clip_grad_norm=0.0, 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) – Loss function for calculating loss

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

  • 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
  • loss (torch.Tensor) – output tensor

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

calc_loss(*args, **kwargs)

compute the loss value :return: the loss value :rtype: torch.Tensor

train()

Sets the model to training mode.

eval()

Sets the model to evaluation mode.