PyTorch模型适配昇腾NPU：从零开始的端到端流程

本文详细介绍了将PyTorch模型部署到昇腾NPU进行推理的全流程。主要内容包括：1）环境搭建，需安装CANN toolkit和torch-npu扩展；2）模型权重转换，将FP32权重转为FP16格式以适配NPU；3）算子对齐检查，处理不支持的算子问题。文章提供了具体代码示例，涵盖ResNet-50和LLaMA-7B等模型的转换方法，并给出三种解决算子不支持问题的方案。通过实战指导帮助开发者顺利完

2502_94507659

204人浏览 · 2026-05-23 18:43:59

2502_94507659 · 2026-05-23 18:43:59 发布

前言

把PyTorch训练的模型跑到昇腾NPU上做推理，看似简单，实际坑很多。环境要搭、权重要转、推理要调、性能要优。这篇文章从零开始，手把手带你走完整个流程。全程实战，不绕弯子。

第一步：环境搭建

要把PyTorch模型跑到昇腾NPU上，需要这些东西：

昇腾NPU硬件：训练用Ascend 910，推理用Ascend 310P。本地开发如果没有硬件，可以用昇腾提供的模拟器（CPU模式）。
CANN toolkit：昇腾的开发者工具包，包含Ascend C编译器、运行时库、调试工具。去昇腾官网下载对应版本（推荐8.0.RC1或以上）。
torch-npu：PyTorch的昇腾后端扩展。它把PyTorch的算子调用桥接到CANN的算子库上。
torchvision/torchaudio：如果模型用了视觉或语音预处理，也需要对应版本。
Python 3.8+：torch-npu支持Python 3.8到3.11。

搭环境的具体步骤（Ubuntu 22.04为例）：

# 1. 安装CANN toolkit
# 去昇腾官网下载CANN 8.0.RC1的toolkit包
# 假设下载到了 ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run
chmod +x ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run
sudo ~/Downloads/Ascend-cann-toolkit_8.0.RC1_linux-x86_64.run --install

# 2. 设置环境变量
# 把这几行加到 ~/.bashrc 里
export ASCEND_HOME=/usr/local/Ascend
export PATH=$ASCEND_HOME/bin:$PATH
export LD_LIBRARY_PATH=$ASCEND_HOME/lib64:$LD_LIBRARY_PATH
export PYTHONPATH=$ASCEND_HOME/python/site-packages:$PYTHONPATH

# 3. 验证CANN安装
ascend-info --version
# 应该输出：CANN 8.0.RC1

# 4. 安装torch-npu
# 先装PyTorch（CPU版本就行，torch-npu会覆盖后端）
pip3 install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cpu

# 再装torch-npu（对应PyTorch 2.1.0和CANN 8.0.RC1）
pip3 install torch-npu==2.1.0rc1 -f https://roma-parent-open.obs.cn-north-4.myhuaweicloud.com/ascend-pytorch/index.html

# 5. 验证torch-npu安装
python3 -c "import torch_npu; print(torch_npu.__version__)"
# 应该输出：2.1.0rc1

# 6. 测试NPU是否可用
python3 -c "
import torch
import torch_npu
x = torch.randn(3, 3).npu()
print(x.device)  # 应该输出：npu:0
"

环境搭好后，可以开始适配模型了。

第二步：模型权重转换

PyTorch模型训练完会保存成.pth或.pt文件，这些文件是CPU/GPU格式的。要跑到昇腾NPU上，需要转换成昇腾格式（主要是FP16量化，节省显存和带宽）。

用一个简单的ResNet-50模型做例子，看怎么转换权重：

# convert_resnet50.py - 转换ResNet-50权重到昇腾格式
import torch
import torch_npu

# 1. 加载PyTorch格式的权重
# 用torchvision提供的预训练权重
from torchvision import models

model = models.resnet50(pretrained=True)
state_dict = model.state_dict()

# 2. 转换成FP16（昇腾NPU的FP16性能最好）
# 为什么要转FP16？因为昇腾910的FP16算力是FP32的2倍
# 而且FP16占的显存只有FP32的一半，能放更大的batch
for key in state_dict:
    state_dict[key] = state_dict[key].to(torch.float16)

# 3. 保存成昇腾格式
# 昇腾格式的权重就是把FP16的state_dict存成.pth文件
# 加载的时候用 torch.load(..., map_location="npu")
torch.save(state_dict, "./resnet50_fp16.pth")

print("权重转换完成！")
print(f"原始模型大小（FP32）：{sum(p.numel() * 4 for p in model.parameters()) / 1024 / 1024:.2f} MB")
print(f"转换后大小（FP16）：{sum(p.numel() * 2 for p in model.parameters()) / 1024 / 1024:.2f} MB")
# 输出：
# 权重转换完成！
# 原始模型大小（FP32）：97.28 MB
# 转换后大小（FP16）：48.64 MB

如果是大模型（比如LLaMA-7B），权重文件很大（13GB+），一次性加载会OOM。需要分片转换：

# convert_llama.py - 转换LLaMA-7B权重（分片处理）
import torch
import os

def convert_llama_sharded(model_path, output_path, dtype=torch.float16):
    """
    分片转换LLaMA权重
    为什么要分片？因为LLaMA-7B的权重有13GB+，一次性加载会OOM
    分片后每个文件只存几层，内存占用可控
    """
    # 1. 列出所有权重文件
    # HuggingFace格式的LLaMA权重是分片的：pytorch_model_00001-of-00002.bin ...
    weight_files = sorted([f for f in os.listdir(model_path) if f.endswith(".bin")])
    
    # 2. 逐个文件转换
    for i, weight_file in enumerate(weight_files):
        print(f"正在转换 {weight_file} ({i+1}/{len(weight_files)})...")
        
        # 加载当前分片
        state_dict = torch.load(os.path.join(model_path, weight_file), map_location="cpu")
        
        # 转换成目标精度
        for key in state_dict:
            state_dict[key] = state_dict[key].to(dtype)
        
        # 保存成昇腾格式
        output_file = weight_file.replace(".bin", f"_{dtype}.pth")
        torch.save(state_dict, os.path.join(output_path, output_file))
        
        # 释放内存
        del state_dict
        torch.cuda.empty_cache()  # 如果用GPU转换的话
        
        print(f"✓ {weight_file} 转换完成")
    
    print("全部转换完成！")

# 用法
convert_llama_sharded(
    model_path="./llama-7b-hf/",
    output_path="./llama-7b-npu/",
    dtype=torch.float16
)

第三步：模型适配（算子对齐）

权重转换完，还需要检查模型里用到的算子，昇腾NPU是否都支持。PyTorch的算子很多，昇腾CANN的算子库只实现了常用的那部分。如果模型里用了不支持的算子，推理会报错。

检查方法很简单：把模型加载到NPU上，跑一次前向推理，看有没有报错。

# check_operator_support.py - 检查模型算子支持情况
import torch
import torch_npu
from torchvision import models

# 1. 加载模型并转到NPU
model = models.resnet50(pretrained=False)
model.load_state_dict(torch.load("./resnet50_fp16.pth"))
model = model.npu()  # 转到NPU上
model.eval()           # 推理模式

# 2. 构造输入并跑前向
input = torch.randn(1, 3, 224, 224, dtype=torch.float16).npu()

try:
    with torch.no_grad():  # 不计算梯度，省显存
        output = model(input)
    print("✓ 所有算子都支持！")
    print(f"输出形状：{output.shape}")
except RuntimeError as e:
    print(f"✗ 有算子不支持：{e}")
    # 错误信息会告诉你哪个算子不支持
    # 需要根据错误信息改模型（换支持的算子）或自定义算子

如果碰到不支持的算子，有三种解决办法：

办法一：换支持的算子。比如模型里用了torch.nn.functional.glu，但昇腾NPU不支持。可以换成torch.nn.functional.silu（昇腾支持）。

办法二：自定义算子。如果实在找不到替代的算子，可以用Ascend C写一个自定义算子，然后注册到torch-npu里。

办法三：用CPU算子。把不支持的算子放到CPU上跑，其他算子继续在NPU上跑。这种做法性能会差一些（因为要频繁在CPU和NPU之间搬运数据），但能跑通。

# fallback_to_cpu.py - 把不支持的算子放到CPU上跑
import torch
import torch_npu

class MixedModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(1024, 4096)
        self.act = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(4096, 1024)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        # 假设 self.fc2 在NPU上不支持，就转到CPU上跑
        x_cpu = x.cpu()
        x_cpu = self.fc2(x_cpu)
        # 再转回NPU
        x = x_cpu.npu()
        return x

model = MixedModel()
# fc1和act在NPU上，fc2在CPU上
# 这种做法能跑通，但性能会差一些

第四步：推理性能调优

模型跑通后，还需要调优性能。昇腾NPU的性能调优主要从这几点入手：

1. 算子融合。把多个小算子合并成一个大算子，减少HBM读写和kernel Launch开销。torch-npu会自动做一部分融合，但有些融合需要手动开启。

# 开启算子融合（torch-npu 2.1.0+）
import torch_npu
torch_npu.enable_fusion(True)  # 开启算子融合

# 或者用环境变量
# export TORCH_NPU_FUSION=1

2. 批量推理。单次推理的延迟可能很高，批量推理能摊薄这个开销。把多个输入拼成一个batch，一次推理搞定。

# batch_inference.py - 批量推理示例
import torch
import torch_npu

model = ...  # 加载模型
model = model.npu()
model.eval()

# 构造batch输入
# 假设单次输入是 (1, 3, 224, 224)，拼成batch=8
inputs = torch.randn(8, 3, 224, 224, dtype=torch.float16).npu()

# 批量推理
with torch.no_grad():
    outputs = model(inputs)

print(outputs.shape)  # (8, 1000)

3. 精度调整。FP16推理的速度比FP32快一倍，但可能有精度损失。如果精度损失太大，可以试试混合精度（有些层FP32，有些层FP16）。

# mixed_precision.py - 混合精度推理
import torch
import torch_npu

model = ...  # 加载模型

# 把某些层转成FP32（精度敏感层）
model.fc1 = model.fc1.to(torch.float32)
model.fc2 = model.fc2.to(torch.float32)

# 其他层保持FP16（计算密集层）
model.conv1 = model.conv1.to(torch.float16)
model.conv2 = model.conv2.to(torch.float16)

# 推理时要手动处理精度转换
input = torch.randn(1, 3, 224, 224, dtype=torch.float16).npu()
with torch.no_grad():
    # conv层是FP16，input也是FP16
    x = model.conv1(input)
    x = model.conv2(x)
    # fc层是FP32，需要转成FP32
    x = x.to(torch.float32)
    x = model.fc1(x)
    x = model.fc2(x)
    # 转回FP16（如果后面还有FP16层的话）
    x = x.to(torch.float16)

4. KV-Cache（大模型专用）。推理大模型时，每生成一个token都要重新算历史token的Key和Value，这很浪费。KV-Cache把历史token的K和V缓存起来，新token只需要算自己的K和V。

# kv_cache.py - KV-Cache示例（简化版）
import torch
import torch_npu

class TransformerLayerWithKVCache(torch.nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        self.self_attn = torch.nn.MultiheadAttention(hidden_size, num_heads)
        self.kv_cache = None  # KV-Cache
        
    def forward(self, x):
        # x shape: (batch, seq_len, hidden_size)
        if self.kv_cache is None:
            # 第一次forward，没有缓存，正常算
            attn_output, _ = self.self_attn(x, x, x)
            # 保存K和V到缓存
            # 这里简化了，实际要实现MultiheadAttention的K/V分离
            self.kv_cache = ...
        else:
            # 有缓存，只需要算新token的K和V
            new_k, new_v = ...  # 算新token的K和V
            # 拼到缓存里
            # 这样历史token的K和V不用重新算
            ...
        return attn_output

第五步：部署推理服务

模型调优完，最后一步是部署成推理服务。最简单的做法是用Flask或FastAPI包一层HTTP接口。

# deploy_resnet50.py - 部署ResNet-50推理服务
from flask import Flask, request, jsonify
import torch
import torch_npu
from torchvision import models, transforms
from PIL import Image

app = Flask(__name__)

# 1. 加载模型
model = models.resnet50(pretrained=False)
model.load_state_dict(torch.load("./resnet50_fp16.pth"))
model = model.npu()
model.eval()

# 2. 定义预处理pipeline
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                        std=[0.229, 0.224, 0.225])
])

# 3. 定义推理接口
@app.route("/predict", methods=["POST"])
def predict():
    # 读取上传的图片
    file = request.files["image"]
    img = Image.open(file.stream)
    
    # 预处理
    img_t = preprocess(img)
    img_t = img_t.unsqueeze(0).npu()  # 加batch维度，转到NPU
    
    # 推理
    with torch.no_grad():
        output = model(img_t)
    
    # 后处理：取top-5
    prob = torch.softmax(output[0], dim=0)
    top5_prob, top5_idx = torch.topk(prob, 5)
    
    # 转成Python list（方便JSON序列化）
    result = [
        {"class_id": int(top5_idx[i]), "probability": float(top5_prob[i])}
        for i in range(5)
    ]
    
    return jsonify({"top5": result})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)

启动服务后，可以用curl测试：

curl -X POST http://localhost:8000/predict \
  -F "image=@./examples/dog.jpg"

# 输出：
# {"top5": [
#   {"class_id": 208, "probability": 0.92},
#   {"class_id": 160, "probability": 0.03},
#   ...
# ]}