torch_npu.optim.NpuFusedSGD

接口原型

CLASS torch_npu.optim.NpuFusedSGD(params, lr, momentum=0, dampening=0, weight_decay=0, nesterov=False)

功能描述

通过张量融合实现的高性能SGD优化器，核心功能和torch.optim.SGD兼容。

SGD的功能和原理可参考https://pytorch.org/docs/2.1/generated/torch.optim.SGD.html#sgd。

参数说明

输入说明

params为参数的可迭代对象或参数组的dict类型。

输出说明

类型为“NpuFusedSGD”的对象。

异常说明

约束说明

NpuFusedSGD的实现机制要求params中的每一个模型参数对象在使用过程中不能被重新申请，否则将导致无法预料的结果。引起模型参数对象被重新申请的操作包括但不限于：

对模型参数对象进行inplace计算，或者读取参数的值，NpuFusedSGD可正常工作。

支持的型号

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import torch
from torch_npu.npu.amp import GradScaler, autocast
from torch_npu.optim import NpuFusedSGD

def _create_simple_params_and_grads():
    params = [
        torch.arange(6).reshape(2, 3).float().npu(),
        torch.arange(12).reshape(4, 3).float().npu(),
        torch.arange(6).reshape(2, 3).half().npu(),
        torch.arange(12).reshape(4, 3).half().npu(),
        torch.arange(15).reshape(5, 3).float().npu(),
        torch.arange(18).reshape(6, 3).half().npu(),
        torch.arange(6).reshape(2, 3).float().npu()
    ]

    for i, p in enumerate(params):
        if i < len(params) - 1:
            p.requires_grad = True
            p.grad = p.clone().detach() / 100.

    return params

opt_kwargs = dict(lr=0.01, momentum=0.9, weight_decay=0.001)
params = _create_simple_params_and_grads()
fused_opt = NpuFusedSGD(params, **opt_kwargs)
with torch.no_grad():
    fused_opt.step()
父主题： torch_npu.optim