当前位置：首页 > news >正文

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

news 2026/5/25 5:20:32

前言

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

环境准备

安装驱动和工具包

# 检查 NPU 设备lspci|grepd802# 安装驱动（已安装则跳过）./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 版本即可）pipinstalltorch==2.1.0+cpu-fhttps://download.pytorch.org/whl/torch_stable.html# 安装 torch_npupipinstalltorch-npu==2.1.0.post3# 验证安装python-c"import torch; import torch_npu; print(torch.npu.is_available())"# 输出：True

模型迁移：三种方式

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

importtorchimporttorch_npu# 自动把模型转到 NPUmodel=MyModel()model=model.npu()# 等价于 model.to("npu:0")# 自动把数据转到 NPUinput_tensor=torch.randn(1,3,224,224).npu()# 推理output=model(input_tensor)print(output.shape)

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

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

importtorchimporttorch.nnasnnimporttorch_npuclassMyModel(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()defforward(self,x):x=self.conv(x)x=self.bn(x)x=self.relu(x)returnx# 手动迁移到 NPUmodel=MyModel().npu()# 检查所有参数都在 NPU 上forname,paraminmodel.named_parameters():print(f"{name}:{param.device}")

输出：

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

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

importtorchimporttorch.onnx# 1. 在 GPU 上加载模型model=MyModel().cuda().eval()# 2. 导出 ONNXdummy_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，模型在 NPUmodel=model.npu()input_tensor=torch.randn(1,3,224,224)# 在 CPU 上output=model(input_tensor)# 报错！# 正确：输入也要转到 NPUinput_tensor=input_tensor.npu()output=model(input_tensor)

问题三：优化器状态

# 错误：优化器在 CPUmodel=model.npu()optimizer=torch.optim.Adam(model.parameters())# 优化器在 CPU# 正确：优化器参数也要在 NPUmodel=model.npu()optimizer=torch.optim.Adam(model.parameters())# PyTorch 会自动处理# 手动检查forparaminoptimizer.param_groups[0]['params']:print(param.device)# 应该是 npu:0

精度验证

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

逐层对比

importtorchimportnumpyasnpdefcompare_model_outputs(gpu_model,npu_model,input_tensor):"""对比 GPU 和 NPU 模型的输出"""# GPU 推理gpu_model.eval()withtorch.no_grad():gpu_output=gpu_model(input_tensor.cuda())# NPU 推理npu_model.eval()withtorch.no_grad():npu_input=input_tensor.npu()npu_output=npu_model(npu_input).cpu()# 转成 NumPygpu_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}")# 判断是否通过ifcosine_sim>0.99andmax_abs_error<0.01:print("✅ 精度验证通过")returnTrueelse:print("❌ 精度验证失败")returnFalse# 使用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)

完整验证流程

deffull_precision_validation(gpu_model,npu_model,dataloader):"""用完整验证集做精度验证"""gpu_model.eval()npu_model.eval()gpu_outputs=[]npu_outputs=[]withtorch.no_grad():forinput_tensor,_indataloader:# GPUgpu_out=gpu_model(input_tensor.cuda()).cpu()gpu_outputs.append(gpu_out)# NPUnpu_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}")returncosine_sim>0.99

性能优化

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

开启混合精度

fromtorch_npu.contribimporttransfer_to_npu# 开启自动混合精度model=transfer_to_npu(model)# 或者用 AMPscaler=torch.npu.amp.GradScaler()forinput_tensor,targetindataloader:withtorch.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

性能对比

importtimedefbenchmark(model,input_tensor,iterations=100):"""性能测试"""# 预热for_inrange(10):model(input_tensor)torch.npu.synchronize()start=time.time()for_inrange(iterations):model(input_tensor)torch.npu.synchronize()end=time.time()latency=(end-start)/iterations*1000# msthroughput=iterations/(end-start)# FPSprint(f"延迟:{latency:.2f}ms")print(f"吞吐量:{throughput:.2f}FPS")returnlatency,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 并编译

# 导出 ONNXmodel.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"}})

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

用 Ascend CL 做推理

#include<acl/acl.h>intmain(){// 初始化aclInit(nullptr);aclrtSetDevice(0);// 加载模型uint32_tmodelId;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;intpred=std::max_element(probs,probs+1000)-probs;std::cout<<"Prediction: "<<pred<<std::endl;// 释放资源aclmdlUnload(modelId);aclrtResetDevice(0);aclFinalize();return0;}

参考资源

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 做高性能推理。

查看全文

http://www.jsqmd.com/news/881591/