PyTorch GPU2Ascend工具可将基于GPU的训练脚本迁移为支持NPU的脚本,大幅度提高脚本迁移速度,降低开发者的工作量。本样例可以让开发者快速体验PyTorch GPU2Ascend工具的迁移效率。原基于GPU的训练脚本迁移成功后可在NPU上运行。
PyTorch GPU2Ascend工具为分析迁移工具之一,更多介绍请参见《分析迁移工具》。
本样例选用ResNet50模型。
从pytorch的examples仓库下载main.py文件,并参考相应README在GPU环境完成训练任务,数据集可以选择README中介绍的Dummy Data。
安装CANN软件后,使用CANN运行用户进行编译、运行时,需要以CANN运行用户登录环境,执行source ${install_path}/set_env.sh命令设置环境变量。其中${install_path}为CANN软件的安装目录,例如:/usr/local/Ascend/ascend-toolkit。
1 2 3 4 5 | 23 24 import torch_npu 25 from torch_npu.contrib import transfer_to_npu 26 # 在24 25行插入的代码为自动迁移的库代码,可以在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) |