DeepSpeedOptimWrapper

class mmengine.optim.DeepSpeedOptimWrapper(optimizer)

backward(loss, **kwargs)

“Perform gradient back propagation.

参数

loss (torch.Tensor) –

返回类型

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.

参数

state_dict (dict) – The state dictionary of optimizer.

返回类型

None

state_dict()

A wrapper of Optimizer.state_dict.

返回类型

dict

step(**kwargs)

Call the step method of optimizer.

update_params(loss)

Update parameters in optimizer.

返回类型

None

zero_grad(**kwargs)

A wrapper of Optimizer.zero_grad.

返回类型

None