DeepSpeedOptimWrapper

class mmengine._strategy.deepspeed.DeepSpeedOptimWrapper(optimizer)

backward(loss, **kwargs)

“Perform gradient back propagation.

Parameters

loss (torch.Tensor) –

Return type

None

load_state_dict(state_dict)

A wrapper of Optimizer.load_state_dict. load the state dict of optimizer.

Provide unified load_state_dict interface compatible with automatic mixed precision training. Subclass can overload this method to implement the required logic. For example, the state dictionary of GradScaler should be loaded when training with torch.cuda.amp.

Parameters

state_dict (dict) – The state dictionary of optimizer.

Return type

None

state_dict()

A wrapper of Optimizer.state_dict.

Return type

dict

step(**kwargs)

Call the step method of optimizer.

update_params(loss)

Update parameters in optimizer.

Return type

None

zero_grad(**kwargs)

A wrapper of Optimizer.zero_grad.

Return type

None