PyTorch 模型迁移实战：从 GPU 到 NPU

passion098

446人浏览 · 2026-05-24 12:55:17

passion098 · 2026-05-24 12:55:17 发布

前言

把在 GPU 上训练好的 PyTorch 模型迁到昇腾 NPU，大部分时候不难，但细节很多。这篇文章讲一个完整的迁移流程，从环境准备到性能验证。

环境准备

安装驱动和工具包

# 检查 NPU 设备
lspci | grep d802

# 安装驱动（已安装则跳过）
./Ascend-driver_24.1.RC3_linux-x86_64.run --install

# 安装 CANN 工具包
./Ascend-cann-toolkit_7.0.RC1_linux-x86_64.run --install

# 设置环境变量
source /usr/local/Ascend/ascend-toolkit/set_env.sh

安装 PyTorch 和 torch_npu

# 安装 PyTorch（CPU 版本即可）
pip install torch==2.1.0+cpu -f https://download.pytorch.org/whl/torch_stable.html

# 安装 torch_npu
pip install torch-npu==2.1.0.post3

# 验证安装
python -c "import torch; import torch_npu; print(torch.npu.is_available())"
# 输出：True

模型迁移：三种方式

方式一：自动迁移（最简单）

import torch
import torch_npu

# 自动把模型转到 NPU
model = MyModel()
model = model.npu()  # 等价于 model.to("npu:0")

# 自动把数据转到 NPU
input_tensor = torch.randn(1, 3, 224, 224).npu()

# 推理
output = model(input_tensor)
print(output.shape)

适合：模型没有自定义算子、没有动态控制流的情况。

方式二：手动迁移（最可靠）

import torch
import torch.nn as nn
import torch_npu

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3, padding=1)
        self.bn = nn.BatchNorm2d(64)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

# 手动迁移到 NPU
model = MyModel().npu()

# 检查所有参数都在 NPU 上
for name, param in model.named_parameters():
    print(f"{name}: {param.device}")

输出：

conv.weight: npu:0
conv.bias: npu:0
bn.weight: npu:0
bn.bias: npu:0

方式三：ONNX 中间表示（最通用）

import torch
import torch.onnx

# 1. 在 GPU 上加载模型
model = MyModel().cuda().eval()

# 2. 导出 ONNX
dummy_input = torch.randn(1, 3, 224, 224).cuda()
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=11,
    input_names=["input"],
    output_names=["output"]
)

# 3. 用 ATC 编译成 .om
!atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --input_shape="input:1,3,224,224"

适合：从其他框架（TensorFlow、MXNet）迁移，或者需要部署到生产环境的情况。

常见迁移问题

问题一：CUDA 算子不支持

# GPU 代码（报错）
output = torch.nn.functional.interpolate(
    input,
    scale_factor=2,
    mode='bilinear',
    align_corners=False
)

# NPU 代码（正常）
output = torch.nn.functional.interpolate(
    input,
    scale_factor=2,
    mode='nearest'  # NPU 对 nearest 支持更好
)

问题二：设备不匹配

# 错误：输入在 CPU，模型在 NPU
model = model.npu()
input_tensor = torch.randn(1, 3, 224, 224)  # 在 CPU 上
output = model(input_tensor)  # 报错！

# 正确：输入也要转到 NPU
input_tensor = input_tensor.npu()
output = model(input_tensor)

问题三：优化器状态

# 错误：优化器在 CPU
model = model.npu()
optimizer = torch.optim.Adam(model.parameters())  # 优化器在 CPU

# 正确：优化器参数也要在 NPU
model = model.npu()
optimizer = torch.optim.Adam(model.parameters())  # PyTorch 会自动处理

# 手动检查
for param in optimizer.param_groups[0]['params']:
    print(param.device)  # 应该是 npu:0

精度验证

迁移后最重要的一步：验证 NPU 上的输出和 GPU 上是否一致。

逐层对比

import torch
import numpy as np

def compare_model_outputs(gpu_model, npu_model, input_tensor):
    """对比 GPU 和 NPU 模型的输出"""
    # GPU 推理
    gpu_model.eval()
    with torch.no_grad():
        gpu_output = gpu_model(input_tensor.cuda())
    
    # NPU 推理
    npu_model.eval()
    with torch.no_grad():
        npu_input = input_tensor.npu()
        npu_output = npu_model(npu_input).cpu()
    
    # 转成 NumPy
    gpu_output_np = gpu_output.cpu().numpy()
    npu_output_np = npu_output.numpy()
    
    # 计算余弦相似度
    cosine_sim = np.dot(gpu_output_np.flatten(), npu_output_np.flatten()) / \
                (np.linalg.norm(gpu_output_np) * np.linalg.norm(npu_output_np))
    
    # 计算最大绝对误差
    max_abs_error = np.max(np.abs(gpu_output_np - npu_output_np))
    
    print(f"余弦相似度: {cosine_sim:.6f}")
    print(f"最大绝对误差: {max_abs_error:.6f}")
    
    # 判断是否通过
    if cosine_sim > 0.99 and max_abs_error < 0.01:
        print("✅ 精度验证通过")
        return True
    else:
        print("❌ 精度验证失败")
        return False

# 使用
gpu_model = torch.load("gpu_model.pth")
npu_model = torch.load("gpu_model.pth").npu()
input_tensor = torch.randn(1, 3, 224, 224)

compare_model_outputs(gpu_model, npu_model, input_tensor)

完整验证流程

def full_precision_validation(gpu_model, npu_model, dataloader):
    """用完整验证集做精度验证"""
    gpu_model.eval()
    npu_model.eval()
    
    gpu_outputs = []
    npu_outputs = []
    
    with torch.no_grad():
        for input_tensor, _ in dataloader:
            # GPU
            gpu_out = gpu_model(input_tensor.cuda()).cpu()
            gpu_outputs.append(gpu_out)
            
            # NPU
            npu_out = npu_model(input_tensor.npu()).cpu()
            npu_outputs.append(npu_out)
    
    # 拼接
    gpu_outputs = torch.cat(gpu_outputs, dim=0).numpy()
    npu_outputs = torch.cat(npu_outputs, dim=0).numpy()
    
    # 指标
    cosine_sim = np.dot(gpu_outputs.flatten(), npu_outputs.flatten()) / \
                (np.linalg.norm(gpu_outputs) * np.linalg.norm(npu_outputs))
    
    max_abs_error = np.max(np.abs(gpu_outputs - npu_outputs))
    mean_abs_error = np.mean(np.abs(gpu_outputs - npu_outputs))
    
    print(f"余弦相似度: {cosine_sim:.6f}")
    print(f"最大绝对误差: {max_abs_error:.6f}")
    print(f"平均绝对误差: {mean_abs_error:.6f}")
    
    return cosine_sim > 0.99

性能优化

迁移完成后，做基本的性能优化。

开启混合精度

from torch_npu.contrib import transfer_to_npu

# 开启自动混合精度
model = transfer_to_npu(model)

# 或者用 AMP
scaler = torch.npu.amp.GradScaler()

for input_tensor, target in dataloader:
    with torch.npu.amp.autocast():
        output = model(input_tensor.npu())
        loss = criterion(output, target.npu())
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

用 AOE 做自动调优

# 调优模型
aoe --model=model.onnx \
    --framework=5 \
    --job_type=2 \
    --mode=rl \
    --output=model_optimized

性能对比

import time

def benchmark(model, input_tensor, iterations=100):
    """性能测试"""
    # 预热
    for _ in range(10):
        model(input_tensor)
    
    torch.npu.synchronize()
    start = time.time()
    
    for _ in range(iterations):
        model(input_tensor)
    
    torch.npu.synchronize()
    end = time.time()
    
    latency = (end - start) / iterations * 1000  # ms
    throughput = iterations / (end - start)  # FPS
    
    print(f"延迟: {latency:.2f} ms")
    print(f"吞吐量: {throughput:.2f} FPS")
    
    return latency, throughput

# 对比
gpu_latency, gpu_fps = benchmark(gpu_model, input_tensor.cuda())
npu_latency, npu_fps = benchmark(npu_model, input_tensor.npu())

print(f"NPU/GPU 延迟比: {npu_latency / gpu_latency:.2f}x")
print(f"NPU/GPU 吞吐量比: {npu_fps / gpu_fps:.2f}x")

部署到生产环境

导出 ONNX 并编译

# 导出 ONNX
model.eval()
dummy_input = torch.randn(1, 3, 224, 224).npu()
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=11,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}}
)

# 编译成 .om
atc --model=model.onnx \
    --framework=5 \
    --output=model \
    --enable_fusion=true \
    --op_precision_mode=allow_fp32_to_fp16

用 Ascend CL 做推理

#include <acl/acl.h>

int main() {
    // 初始化
    aclInit(nullptr);
    aclrtSetDevice(0);
    
    // 加载模型
    uint32_t modelId;
    aclmdlLoadFromFile("model.om", &modelId);
    
    // 准备输入
    aclmdlDataset* input = aclmdlCreateDataset();
    void* inputBuffer = /* 输入数据 */;
    aclDataBuffer* inputData = aclCreateDataBuffer(inputBuffer, inputSize);
    aclmdlAddDatasetBuffer(input, inputData);
    
    // 推理
    aclmdlDataset* output = aclmdlCreateDataset();
    aclmdlExecuteAsync(modelId, input, output, nullptr);
    aclrtSynchronizeStream(nullptr);
    
    // 获取输出
    aclDataBuffer* outputData = aclmdlGetDatasetBuffer(output, 0);
    void* outputBuffer = aclGetDataBufferAddr(outputData);
    
    // 后处理
    float* probs = (float*)outputBuffer;
    int pred = std::max_element(probs, probs + 1000) - probs;
    std::cout << "Prediction: " << pred << std::endl;
    
    // 释放资源
    aclmdlUnload(modelId);
    aclrtResetDevice(0);
    aclFinalize();
    
    return 0;
}

参考资源

PyTorch 模型迁移指南: https://www.hiascend.com/document/detail/zh/CANN/
torch_npu API 文档: https://gitee.com/ascend/pytorch
ONNX 模型导出: https://pytorch.org/docs/stable/onnx.html
精度调优指南: https://www.hiascend.com/document/detail/zh/CANN/

总结

PyTorch 模型迁移到 NPU 有三种方式：自动迁移（最简单）、手动迁移（最可靠）、ONNX 中间表示（最通用）。迁移后要做精度验证，确保余弦相似度 > 0.99。性能优化包括开启混合精度、用 AOE 调优、算子融合。生产部署时，导出 ONNX 并编译成 .om，再用 Ascend CL 做高性能推理。