pytorch-adapter：让 PyTorch 模型“无缝”跑在昇腾 NPU 上

之前帮朋友看 PyTorch 模型适配 CANN 的代码，发现他手写了很多适配层——把自己的 MyModel 一层层翻译成 AscendCL 接口，光写适配层就写了 2,000 行。

我告诉他：不用手写，用 pytorch-adapter 就行。 这个适配器能自动把 PyTorch 模型跑在 NPU 上，不用改模型代码（或者只改几行）。

环境准备：装 pytorch-adapter 和依赖

在拆 pytorch-adapter 的用法之前，先把环境装好。不然后面跑代码报“模块找不到”，又得回头查。

第1步：装 CANN（必备）

pytorch-adapter 依赖 CANN 的 AscendCL 接口，得先装 CANN。推荐装 CANN 8.0+（对 PyTorch 2.x 支持更好）。

# 检查 CANN 是否装好
npu-smi info

如果看到 NPU 设备信息，说明 CANN 装好了。

⚠️ 踩坑预警：CANN 版本跟 pytorch-adapter 版本要对应。CANN 8.0 得配 pytorch-adapter v3.x，配错了模型跑不起来。

第2步：装 PyTorch（推荐 2.1+）

pytorch-adapter 支持 PyTorch 1.13-2.1，推荐用 PyTorch 2.1（性能更好）。

# 装 PyTorch 2.1（CPU 版本就行，adapter 会替换成 NPU 版本）
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cpu

⚠️ 踩坑预警：装的是 CPU 版本的 PyTorch（不是 CUDA 版本）。pytorch-adapter 会替换 PyTorch 的后端（从 CUDA 换成 NPU），如果你装了 CUDA 版本，会冲突。

第3步：装 pytorch-adapter

# 方法1：pip 安装（推荐）
pip install pytorch-npu==3.0.0  # 对应 CANN 8.0

# 方法2：源码编译（如果你想改 adapter 的代码）
git clone https://atomgit.com/cann/pytorch-adapter.git
cd pytorch-adapter && git checkout v3.0  # 对应 CANN 8.0
python setup.py install

装完后，Python 里能 import torch 并且 torch.npu.is_available() 返回 True，就说明装好了。

逐步实现：让 PyTorch 模型跑在 NPU 上

第1步：把模型搬到 NPU 上（.npu()）

pytorch-adapter 提供了 .npu() 方法（类似 PyTorch 的 .cuda()），把模型/张量搬到 NPU 上。

import torch
import torch.nn as nn

# 1. 定义模型（跟普通 PyTorch 模型一样）
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 10)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = MyModel()

# 2. 把模型搬到 NPU 上（关键！）
model = model.npu()  # 类似 .cuda()

# 3. 把输入数据搬到 NPU 上
input_data = torch.randn(32, 784).npu()  # 类似 .cuda()

# 4. 前向计算（自动在 NPU 上跑）
output = model(input_data)
print(output.device)  # 输出：npu:0

关键点：

• .npu() 是 pytorch-adapter 自动添加的方法（不用改 PyTorch 源码）
• 模型搬到 NPU 后，所有计算都自动在 NPU 上跑（不用手动调算子）
• ⚠️ 别忘了把输入数据也搬到 NPU 上（input_data = input_data.npu()）。如果模型在 NPU 上，输入数据在 CPU 上，会报“设备不匹配”错误。

第2步：训练模型（自动在 NPU 上算梯度）

pytorch-adapter 自动替换了 PyTorch 的后端（从 CUDA 换成 NPU），所以训练代码不用改（跟 GPU 训练一模一样）。

import torch
import torch.nn as nn
import torch.optim as optim

# 1. 定义模型 + 搬到 NPU
model = MyModel().npu()

# 2. 定义损失函数 + 优化器（跟 GPU 训练一模一样）
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 3. 训练循环（跟 GPU 训练一模一样）
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 把数据搬到 NPU
        data, target = data.npu(), target.npu()
        
        # 前向计算
        output = model(data)
        loss = criterion(output, target)
        
        # 反向传播（自动在 NPU 上算梯度）
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch}, Batch {batch_idx}, Loss {loss.item()}')

关键点：

• loss.backward() 会自动在 NPU 上算梯度（pytorch-adapter 替换了 autograd 后端）
• 优化器会自动更新 NPU 上的参数（不用手动拷贝）
• ⚠️ 如果你用了自定义的 autograd 函数（torch.autograd.Function），得手动改（adapter 没法自动替换）。得把 ctx.cuda() 改成 ctx.npu()。

第3步：保存/加载模型（NPU 格式）

训练完的模型，可以用 PyTorch 标准的保存/加载接口（adapter 自动处理了 NPU 格式）。

# 1. 保存模型（跟 GPU 训练一模一样）
torch.save(model.state_dict(), 'model_npu.pth')

# 2. 加载模型（得先搬到 NPU 上）
model = MyModel().npu()  # 先搬到 NPU
model.load_state_dict(torch.load('model_npu.pth'))
model.eval()  # 推理模式

# 3. 推理
input_data = torch.randn(1, 784).npu()
output = model(input_data)
pred = output.argmax(dim=1)
print(f'Prediction: {pred.item()}')

⚠️ 踩坑预警：保存的模型是 NPU 格式的（不是 CPU 格式）。如果你要拿到 CPU 上推理，得先搬到 CPU（model = model.cpu()），再保存。

性能调优：让模型跑得更快

调优点1：开启算子融合（Graph Mode）

pytorch-adapter 支持 Graph Mode（图模式），会自动融合算子（类似 GE 图引擎的功能），性能提升 20%-30%。

import torch

# 开启 Graph Mode（在模型推理前调用）
torch.npu.enable_graph_mode(True)

# 推理（会自动做算子融合）
input_data = torch.randn(1, 784).npu()
output = model(input_data)

⚠️ 踩坑预警：Graph Mode 不是所有模型都支持。如果你的模型有动态控制流（if/for/while），Graph Mode 会报错。这时候得关掉 Graph Mode（torch.npu.enable_graph_mode(False)）。

调优点2：用 NPU 的加速库（ATB）

如果你用的是 Transformer 类模型（BERT/GPT/LLaMA），可以用 ATB（Ascend Transformer Boost）加速，性能提升 50%-100%。

from atb import TransformerInfer

# 用 ATB 包装模型
model = MyGPTModel().npu()
infer = TransformerInfer(model, max_seq_len=2048, soc_type="Ascend910")

# 推理（自动用 ATB 加速）
input_ids = torch.tensor([[1, 2, 3, 4, 5]]).npu()
output = infer.forward(input_ids)

调优点3：用混合精度训练（FP16）

NPU 的 FP16 算力是 FP32 的 2-3 倍。用混合精度训练，性能提升 50%-80%。

import torch

# 开启自动混合精度（跟 GPU 的 amp 用法一样）
with torch.npu.amp.autocast():
    output = model(input_data)
    loss = criterion(output, target)
    loss.backward()

踩坑实录

踩坑1：模型搬不到 NPU 上（.npu() 报错）

原因：PyTorch 版本跟 pytorch-adapter 不匹配。

解决方案：检查版本对应关系：

PyTorch 版本	pytorch-adapter 版本
1.13	v2.0.x
2.0	v3.0.x
2.1	v3.1.x

踩坑2：训练时梯度算不出来（loss.backward() 报错）

原因：用了自定义的 autograd 函数（torch.autograd.Function），adapter 没法自动替换。

解决方案：手动改自定义函数的代码，把 ctx.cuda() 改成 ctx.npu()。

踩坑3：推理性能不达标（比 GPU 慢）

原因：没开 Graph Mode 或者没用 ATB 加速库。

解决方案：

• 开启 Graph Mode（torch.npu.enable_graph_mode(True)）
• 如果是 Transformer 模型，用 ATB 加速（TransformerInfer）

实践指引

1. 读 pytorch-adapter 源码：从 torch/npu/__init__.py 看起，理解 .npu() 是怎么注入的
2. 跑 pytorch-adapter 的示例：pytorch-adapter 仓库里有现成的示例（examples/ 目录）
3. 调 Graph Mode：如果你的模型推理性能不达标，试试开 Graph Mode
4. 用 ATB 加速 Transformer 模型：如果是 Transformer 类模型，用 ATB 加速，性能提升 50%-100%

仓库链接：
https://atomgit.com/cann/pytorch-adapter
https://atomgit.com/cann/ascend-transformer-boost
https://atomgit.com/cann/runtime