基于华为自研NPU Ascend 910的TensorFlow 1.x训练脚本迁移和使能混合精度记录

简介

基于 TesorFlow 1.x 以 Sess.run 形式搭建入门级——手写数字分类网络，并迁移到 华为自研NPU Ascend 910，同时使能混合精度。

硬件介绍

华为自研 NPU Ascend 910，即昇腾 910 AI 处理器（简称NPU），根据官方介绍，是在2019年发布的人工智能（AI）专用的神经网络处理器，其算力高达256T，最新款算力高达310T，是业界主流芯片算力的2倍。当前业界大多数训练脚本基于 TensorFlow 的 Python API 开发，默认运行在 CPU/GPU/TPU 上，为了使其能够利用昇腾910 AI处理器的澎湃算力执行训练，提升训练性能，需要对训练网络脚本进行简单的迁移适配工作。当前昇腾910 AI处理器上支持 TensorFlow 1.15 的三种 API 开发的训练脚本迁移：分别是Estimator，Sess.run，Keras。 这里以一个Sess.run的手写数字分类网络为例，介绍如何迁移TensorFlow 1.15训练脚本，以支持NPU训练。

同时，华为还推出了自研 AI 计算框架 MindSpore，充分释放 Ascend 910 的澎湃算力，构建全栈全场景 AI 解决方案，有兴趣的同学可以试试。此外，还有Ascend 310，是一款高能效、灵活可编程的人工智能处理器，在典型配置下可以输出 16TOPS@INT8, 8TOPS@FP16，功耗仅为8W。采用自研华为达芬奇架构，集成丰富的计算单元, 提高 AI 计算完备度和效率，进而扩展该芯片的适用性。全 AI 业务流程加速,大幅提高AI全系统的性能，有效降低部署成本。

搭建网络

这里，我们使用华为云 ModelArts 提供的免费资源：Ascend：1*Ascend 910 规格。

下载训练数据集

直接执行以下命令下载手写数字数据集，这个数据集非常经典，可以很方便地下载到。

!wget -N -P /home/ma-user/work/Data https://modelarts-train-ae.obs.cn-north-4.myhuaweicloud.com/train/Data/t10k-images.idx3-ubyte
!wget -N -P /home/ma-user/work/Data https://modelarts-train-ae.obs.cn-north-4.myhuaweicloud.com/train/Data/t10k-labels.idx1-ubyte
!wget -N -P /home/ma-user/work/Data https://modelarts-train-ae.obs.cn-north-4.myhuaweicloud.com/train/Data/train-labels.idx1-ubyte
!wget -N -P /home/ma-user/work/Data https://modelarts-train-ae.obs.cn-north-4.myhuaweicloud.com/train/Data/train-images.idx3-ubyte

导入库文件

要使基于 TensorFlow 开发的训练脚本在昇腾910 AI处理器上训练，需要借助 Tensorflow 框架适配插件（即 TF Adapter），TF Adapter 中提供了适配 Tensorflow 框架的用户 Python 接口，用于 CANN 软件与 TensorFlow 框架对接。因此在训练之前，需要在训练代码中增加：

from npu_bridge.npu_init import *

导入相关库文件。具体代码：

import tensorflow as tf
import numpy as np
import struct
import os
import time 
from npu_bridge.npu_init import *

处理MNIST数据集

此部分代码应该很常见了，和在 GPU 或在 CPU 上运行基本一致，除为了特定需求或性能优化外，一般无需改造。

首先是加载数据集，原始数据为了节省保存空间，也为了加载的快速，一般保存为二进制文件。