文档HOT
下载
中文
注册

模型开发&迁移

PyTorch GPU2Ascend工具可将基于GPU的训练脚本迁移为支持NPU的脚本，大幅度提高脚本迁移速度，降低开发者的工作量。本样例可以让开发者快速体验PyTorch GPU2Ascend工具的迁移效率。原基于GPU的训练脚本迁移成功后可在NPU上运行。

PyTorch GPU2Ascend工具为分析迁移工具之一，更多介绍请参见《分析迁移工具》。

前提条件

本样例选用resnet50模型。

从pytorch的examples仓库下载main.py文件，并参考相应README在GPU环境完成训练任务，数据集可以选择README中介绍的Dummy Data。

环境准备

  1. 准备一台基于Atlas 训练系列产品的训练服务器，并安装NPU驱动固件。
  2. 安装开发套件包Ascend-cann-toolkit，具体请参考《CANN 软件安装指南》的“安装CANN软件包”节点。
  3. 以安装PyTorch 2.1.0版本为例，具体操作请参考《Ascend Extension for PyTorch 配置与安装》。
  4. 配置环境变量。

    安装CANN软件后，使用CANN运行用户进行编译、运行时，需要以CANN运行用户登录环境，执行source ${install_path}/set_env.sh命令设置环境变量。其中${install_path}为CANN软件的安装目录，例如：/usr/local/Ascend/ascend-toolkit。

  5. 下载main.py文件，并上传至训练服务器的个人目录下。

执行迁移

  1. 在训练脚本（main.py文件）中导入自动迁移的库代码。
    1
2
3
4
5
     23
 24 import torch_npu
 25 from torch_npu.contrib import transfer_to_npu
 26
# 在24 25行插入的代码为自动迁移的库代码，可以在NPU环境下直接执行训练
  2. 迁移完成后的训练脚本可在NPU上运行，执行以下训练命令。
    python main.py -a resnet50 -b 32 --gpu 1 --dummy

    如果训练正常进行，开始打印迭代日志，说明训练功能迁移成功，如下所示。

    1
2
3
4
5
6
    Use GPU: 1 for training
=> creating model 'resnet50'
=> Dummy data is used!
Epoch: [0][    1/40037] Time  8.287 ( 8.287)    Data  0.504 ( 0.504)    Loss 7.0919e+00 (7.0919e+00)    Acc@1   0.00 (  0.00)   Acc@5   0.00 (  0.00)
Epoch: [0][   11/40037] Time  0.097 ( 1.268)    Data  0.000 ( 0.479)    Loss 1.5627e+01 (1.8089e+01)    Acc@1   0.00 (  0.00)   Acc@5   3.12 (  0.57)
Epoch: [0][   21/40037] Time  0.096 ( 0.710)    Data  0.000 ( 0.253)    Loss 7.7462e+00 (1.4883e+01)    Acc@1   0.00 (  0.00)   Acc@5   0.00 (  0.60)
  3. 成功保存权重，说明保存权重功能迁移成功。