GPU训练脚本迁移昇腾NPU只需5步简单修改：TorchNPU模型迁移实战指南-育师

GPU训练脚本迁移昇腾NPU只需5步简单修改：TorchNPU模型迁移实战指南

【免费下载链接】pytorch作为 Ascend for PyTorch 社区的核心组件，TorchNPU 是昇腾专为 PyTorch 打造的深度学习适配插件，使 PyTorch 框架能够直接调用昇腾 NPU，为开发者提供昇腾 AI 处理器的超强算力。项目地址: https://gitcode.com/Ascend/pytorch

TorchNPU是昇腾专为 PyTorch 打造的深度学习适配插件，让 PyTorch 框架直接调用昇腾 NPU 的超强算力。本文将带你用5 步简单修改，把一个跑在 GPU 上的 PyTorch 训练脚本迁移到昇腾 NPU 上训练——全程只需改动几行设备相关代码，原脚本结构几乎不用动。

💡 核心思路：TorchNPU 完全继承 PyTorch 的 API 生态，迁移的本质只是把「计算设备」从cuda换成npu，并用配套的自动迁移工具 + AMP 混合精度做少量增强。

上图展示了 TorchNPU 的分层结构：你的训练脚本（PyTorch 三方库）只需通过「前端 API」接入，算子适配、分布式、图模式、DFX 等能力由 Ascend Extension for PyTorch 层统一承接，底层对接 CANN 软件栈。

迁移前准备：TorchNPU 环境一键安装步骤

第 1 步：安装 CANN、PyTorch 与 TorchNPU

迁移的第一前提是把「CANN 软件 + PyTorch + TorchNPU 插件」三件套装齐，支持物理机、虚拟机和容器三类部署场景：

以安装 TorchNPU 2.12.0 为例，核心命令如下：

# 安装 PyTorch（CPU 版本即可，NPU 由 TorchNPU 提供） pip install torch==2.12.0 --index-url https://download.pytorch.org/whl/cpu # 安装 TorchNPU 昇腾NPU适配插件 pip install torch-npu==2.12.0

然后初始化 CANN 环境：

source /usr/local/Ascend/cann/set_env.sh

⚠️版本匹配是关键：TorchNPU / PyTorch / CANN / Python 四者必须按兼容矩阵搭配，具体版本对照见 COMPATIBILITY.md，支持的 PyTorch 版本声明在 version.txt。完整安装细节可参考官方文档目录 docs/zh/installation_guide/。

开启自动迁移：transfer_to_npu 两行搞定设备切换

第 2 步：用 torch_npu.contrib 自动迁移

对于「脚本里显式写了cuda设备」的场景，TorchNPU 提供了自动迁移工具，只需在脚本头部加两行代码：

import torch_npu from torch_npu.contrib import transfer_to_npu # 开启自动迁移

这个工具会自动把脚本中的 GPU 设备操作改写为 NPU 设备操作。其源码位于 torch_npu/contrib/transfer_to_npu.py，对应的功能验证用例在 test/contrib/test_transfer_to_npu.py。

🎯适合谁用：训练脚本结构简单、设备声明集中、没有复杂自定义后端的场景。这也是官方快速入门样例采用的方式（见 docs/zh/user_guide/quick_start.md 中的「模型迁移训练」章节）。

手工迁移：把 cuda 替换成 npu 的三个位置

第 3 步：手工修改设备（不开自动迁移时的替代方案）

自动迁移不适用时，手工迁移也非常简单——全文找替换，只动设备相关的地方：

同时建议在脚本头部显式导入，确保 NPU 设备正确初始化：

import torch import torch_npu

补充两个实用技巧：

查询可用卡数：torch.npu.device_count()；
确认某个算子是否支持：可参考插件内算子接口清单 torch_npu/acl.json；
个别算子 NPU 暂不支持时，可开启 CPU 回退机制兜底，相关机制见 test/npu/test_cpu_fallback_control.py。

开启 AMP 混合精度：NPU 训练的加速关键

第 4 步：配置 AMP 混合精度训练

昇腾910 系列芯片由于架构特性，在迁移完成、正式训练前必须开启混合精度；A2/A3/950DT 系列则可按需选择。核心改动分三处：

① 导入 AMP 模块：

from torch_npu.npu import amp

② 在模型和优化器定义之后，创建 GradScaler：

scaler = amp.GradScaler()

③ 训练循环中用 amp.autocast 包裹前向计算，并用 scaler 管理反向传播：

for imgs, labels in train_dataloader: imgs = imgs.to(device) labels = labels.to(device) with amp.autocast(): outputs = model(imgs) # 前向计算 loss = loss_func(outputs, labels) optimizer.zero_grad() scaler.scale(loss).backward() # loss 缩放并反向传播 scaler.step(optimizer) # 更新参数（自动 unscaling） scaler.update() # 动态更新 loss scaling 系数

注意 AMP 的 GradScaler 来自torch_npu.npu.amp，而不是PyTorch 原生的torch.cuda.amp——这是 NPU 迁移中最容易踩的坑。