训练推理示例

本页给出在昇腾 NPU 上运行 PyTorch 的最小可运行示例，以及基本的调优思路。

1. 最小示例：张量计算

验证 NPU 是否可用的最小代码：

import torch
import torch_npu

a = torch.randn(2, 3).npu()
b = torch.randn(2, 3).npu()
c = a + b
print(c)
print(c.device)        # 预期输出 npu:0

2. 一个简单的训练循环

import torch
import torch_npu
import torch.nn as nn

device = 'npu'

model = nn.Linear(10, 2).to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for step in range(100):
    x = torch.randn(32, 10).to(device)
    y = torch.randint(0, 2, (32,)).to(device)

    optimizer.zero_grad()
    out = model(x)
    loss = criterion(out, y)
    loss.backward()
    optimizer.step()

    if step % 20 == 0:
        print(f'step {step}, loss {loss.item():.4f}')

与标准 PyTorch 训练循环相比，唯一的差别就是 import torch_npu 和 device='npu'。

3. 推理示例

model.eval()
with torch.no_grad():
    x = torch.randn(1, 10).to('npu')
    pred = model(x)
    print(pred.argmax(dim=1))

4. 性能调优思路

• 混合精度：使用 torch.npu.amp 自动混合精度，提升吞吐、降低显存占用。
• 算子下发与同步：注意减少 NPU 与 CPU 间不必要的同步（如频繁 .item() / .cpu()）。
• 多卡并行：用 HCCL 后端做数据并行，充分利用 Atlas 800（9010） 的 8 颗 NPU。

更系统的调优方法见 训练性能调优。

衔接课题

掌握基础用法后，可结合 古蜀道自动识别 课题，在 NPU 上完成一个完整的遥感图像识别任务。