前言把在 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# 安装 PyTorchCPU 版本即可pipinstalltorch2.1.0cpu-fhttps://download.pytorch.org/whl/torch_stable.html# 安装 torch_npupipinstalltorch-npu2.1.0.post3# 验证安装python-cimport torch; import torch_npu; print(torch.npu.is_available())# 输出True模型迁移三种方式方式一自动迁移最简单importtorchimporttorch_npu# 自动把模型转到 NPUmodelMyModel()modelmodel.npu()# 等价于 model.to(npu:0)# 自动把数据转到 NPUinput_tensortorch.randn(1,3,224,224).npu()# 推理outputmodel(input_tensor)print(output.shape)适合模型没有自定义算子、没有动态控制流的情况。方式二手动迁移最可靠importtorchimporttorch.nnasnnimporttorch_npuclassMyModel(nn.Module):def__init__(self):super().__init__()self.convnn.Conv2d(3,64,3,padding1)self.bnnn.BatchNorm2d(64)self.relunn.ReLU()defforward(self,x):xself.conv(x)xself.bn(x)xself.relu(x)returnx# 手动迁移到 NPUmodelMyModel().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 上加载模型modelMyModel().cuda().eval()# 2. 导出 ONNXdummy_inputtorch.randn(1,3,224,224).cuda()torch.onnx.export(model,dummy_input,model.onnx,opset_version11,input_names[input],output_names[output])# 3. 用 ATC 编译成 .om!atc--modelmodel.onnx \--framework5\--outputmodel \--input_shapeinput:1,3,224,224适合从其他框架TensorFlow、MXNet迁移或者需要部署到生产环境的情况。常见迁移问题问题一CUDA 算子不支持# GPU 代码报错outputtorch.nn.functional.interpolate(input,scale_factor2,modebilinear,align_cornersFalse)# NPU 代码正常outputtorch.nn.functional.interpolate(input,scale_factor2,modenearest# NPU 对 nearest 支持更好)问题二设备不匹配# 错误输入在 CPU模型在 NPUmodelmodel.npu()input_tensortorch.randn(1,3,224,224)# 在 CPU 上outputmodel(input_tensor)# 报错# 正确输入也要转到 NPUinput_tensorinput_tensor.npu()outputmodel(input_tensor)问题三优化器状态# 错误优化器在 CPUmodelmodel.npu()optimizertorch.optim.Adam(model.parameters())# 优化器在 CPU# 正确优化器参数也要在 NPUmodelmodel.npu()optimizertorch.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_outputgpu_model(input_tensor.cuda())# NPU 推理npu_model.eval()withtorch.no_grad():npu_inputinput_tensor.npu()npu_outputnpu_model(npu_input).cpu()# 转成 NumPygpu_output_npgpu_output.cpu().numpy()npu_output_npnpu_output.numpy()# 计算余弦相似度cosine_simnp.dot(gpu_output_np.flatten(),npu_output_np.flatten())/\(np.linalg.norm(gpu_output_np)*np.linalg.norm(npu_output_np))# 计算最大绝对误差max_abs_errornp.max(np.abs(gpu_output_np-npu_output_np))print(f余弦相似度:{cosine_sim:.6f})print(f最大绝对误差:{max_abs_error:.6f})# 判断是否通过ifcosine_sim0.99andmax_abs_error0.01:print(✅ 精度验证通过)returnTrueelse:print(❌ 精度验证失败)returnFalse# 使用gpu_modeltorch.load(gpu_model.pth)npu_modeltorch.load(gpu_model.pth).npu()input_tensortorch.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_outgpu_model(input_tensor.cuda()).cpu()gpu_outputs.append(gpu_out)# NPUnpu_outnpu_model(input_tensor.npu()).cpu()npu_outputs.append(npu_out)# 拼接gpu_outputstorch.cat(gpu_outputs,dim0).numpy()npu_outputstorch.cat(npu_outputs,dim0).numpy()# 指标cosine_simnp.dot(gpu_outputs.flatten(),npu_outputs.flatten())/\(np.linalg.norm(gpu_outputs)*np.linalg.norm(npu_outputs))max_abs_errornp.max(np.abs(gpu_outputs-npu_outputs))mean_abs_errornp.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_sim0.99性能优化迁移完成后做基本的性能优化。开启混合精度fromtorch_npu.contribimporttransfer_to_npu# 开启自动混合精度modeltransfer_to_npu(model)# 或者用 AMPscalertorch.npu.amp.GradScaler()forinput_tensor,targetindataloader:withtorch.npu.amp.autocast():outputmodel(input_tensor.npu())losscriterion(output,target.npu())scaler.scale(loss).backward()scaler.step(optimizer)scaler.update()optimizer.zero_grad()用 AOE 做自动调优# 调优模型aoe--modelmodel.onnx\--framework5\--job_type2\--moderl\--outputmodel_optimized性能对比importtimedefbenchmark(model,input_tensor,iterations100):性能测试# 预热for_inrange(10):model(input_tensor)torch.npu.synchronize()starttime.time()for_inrange(iterations):model(input_tensor)torch.npu.synchronize()endtime.time()latency(end-start)/iterations*1000# msthroughputiterations/(end-start)# FPSprint(f延迟:{latency:.2f}ms)print(f吞吐量:{throughput:.2f}FPS)returnlatency,throughput# 对比gpu_latency,gpu_fpsbenchmark(gpu_model,input_tensor.cuda())npu_latency,npu_fpsbenchmark(npu_model,input_tensor.npu())print(fNPU/GPU 延迟比:{npu_latency/gpu_latency:.2f}x)print(fNPU/GPU 吞吐量比:{npu_fps/gpu_fps:.2f}x)部署到生产环境导出 ONNX 并编译# 导出 ONNXmodel.eval()dummy_inputtorch.randn(1,3,224,224).npu()torch.onnx.export(model,dummy_input,model.onnx,opset_version11,input_names[input],output_names[output],dynamic_axes{input:{0:batch}})# 编译成 .omatc--modelmodel.onnx\--framework5\--outputmodel\--enable_fusiontrue\--op_precision_modeallow_fp32_to_fp16用 Ascend CL 做推理#includeacl/acl.hintmain(){// 初始化aclInit(nullptr);aclrtSetDevice(0);// 加载模型uint32_tmodelId;aclmdlLoadFromFile(model.om,modelId);// 准备输入aclmdlDataset*inputaclmdlCreateDataset();void*inputBuffer/* 输入数据 */;aclDataBuffer*inputDataaclCreateDataBuffer(inputBuffer,inputSize);aclmdlAddDatasetBuffer(input,inputData);// 推理aclmdlDataset*outputaclmdlCreateDataset();aclmdlExecuteAsync(modelId,input,output,nullptr);aclrtSynchronizeStream(nullptr);// 获取输出aclDataBuffer*outputDataaclmdlGetDatasetBuffer(output,0);void*outputBufferaclGetDataBufferAddr(outputData);// 后处理float*probs(float*)outputBuffer;intpredstd::max_element(probs,probs1000)-probs;std::coutPrediction: predstd::endl;// 释放资源aclmdlUnload(modelId);aclrtResetDevice(0);aclFinalize();return0;}参考资源PyTorch 模型迁移指南: https://www.hiascend.com/document/detail/zh/CANN/torch_npu API 文档: https://gitee.com/ascend/pytorchONNX 模型导出: 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 做高性能推理。
