把基于PyTorch的模型迁移到华为Ascend910芯片的硬件平台上运行
一 概述:
近几年,随着深度学习的快速发展,相关理论和基础架构得到了很大突破,它们奠定了深度学习繁荣发展的基础。这其中涌现了几个著名的深度学习的平台。
目前国外主流的学习框架有Google的TensorFlow,亚马逊的MXNet,Facebook的Pytorch、Caffe等。
华为也推出了自己的一个全场景深度学习框架MindSpore,旨在实现易开发、高效执行、全场景覆盖三大目标。
CANN是华为昇腾AI基础软硬件平台的核心,以极致性能、极简开发、开放生态为目标,通过CANN可以高效发挥华为各种自研芯片的能力。
从上面的架构图上可以看到,在华为的AI生态上,MindSpore是最佳开发框架。但是MindSport相对比较新,学习需要投入一定的时间。为了减少开发人员的学习成本,华为提供一个解决方案:通过框架适配插件,只要少量修改代码,就可以把基于TensorFlow、PyTorch的代码迁移到CANN运行环境下,然后通过CANN运行在华为的硬件平台上。这样既可以使用熟悉的TensorFlow、PyTorch框架的函数。同时也可以高效发挥华为晟腾的硬件的能力。本文主要介绍的是基于Pytorch的迁移训练。
二 迁移流程:
将基于PyTorch的训练脚本迁移到昇腾AI处理器上进行训练,目前有以下3种方式:自动迁移(推荐)、工具迁移、手工迁移,且迁移前要保证该脚本能在GPU、CPU上运行。
自动迁移:在训练脚本中导入脚本转换库,然后拉起脚本执行训练。训练脚本在运行的同时,会自动将脚本中的CUDA接口替换为昇腾AI处理器支持的NPU接口。整体过程为边训练边转换。
工具迁移:训练前,通过脚本迁移工具,自动将训练脚本中的CUDA接口替换为昇腾AI处理器支持的NPU接口,并生成迁移报告(脚本转换日志、不支持算子的列表、脚本修改记录)。训练时,运行转换后的脚本。整体过程为先转换脚本,再进行训练。
手工迁移:算法工程师通过对模型的分析、GPU与NPU代码的对比进而对训练脚本进行修改,以支持在昇腾AI处理器上执行训练。
三 实际操作步骤:
3.1 建立开发环境
先注册镜像:输入 swr.cn-north-4.myhuaweicloud.com/hehang/ascend_pytorch_1.8.1:v2
然后使用这个镜像创建一个Notebook环境,该环境包含了PyTorch,华为的PyTorch适配包,迁移样例代码。
3.2 使用原生torch训练
3.3 自动迁移,使用华为昇腾910神经网络芯片训练
3.4 用工具迁移,使用华为昇腾910神经网络芯片训练
3.5 手工迁移,使用华为昇腾910神经网络芯片训练
四 总结
华为通过提PyTorch框架适配插件,可以把基于TensorFlow、PyTorch的代码迁移到基于华为CANN-昇腾AI基础软硬件平台上运行。同时提供了云上的镜像包,方便用户学习和使用。