DeepSpeedOptimWrapper — mmengine 0.10.7 documentation (original) (raw)

class mmengine._strategy.deepspeed.DeepSpeedOptimWrapper(optimizer)[source]¶

backward(loss, **kwargs)[source]¶

“Perform gradient back propagation.

Parameters:

loss (Tensor) –

Return type:

None

load_state_dict(state_dict)[source]¶

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

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()[source]¶

A wrapper of Optimizer.state_dict.

Return type:

dict

step(**kwargs)[source]¶

Call the step method of optimizer.

update_params(loss)[source]¶

Update parameters in optimizer.

Return type:

None

zero_grad(**kwargs)[source]¶

A wrapper of Optimizer.zero_grad.

Return type:

None