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

Tianyi_Li 发表于 2022/05/30 19:45:26 2022/05/30
【摘要】 基于 TesorFlow 1.x 以 Sess.run 形式搭建简单网络,并迁移到 华为自研NPU Ascend 910,同时使能混合精度。

简介

基于 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

image.png

导入库文件

要使基于 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 上运行基本一致,除为了特定需求或性能优化外,一般无需改造。

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

# 加载数据集
def load_image_set(filename):
    print ("load image set",filename)
    binfile = open(filename, 'rb')  # 读取二进制文件
    buffers = binfile.read()
    head = struct.unpack_from('>IIII', buffers, 0)  # 读取前四个整数,返回一个元组
    offset = struct.calcsize('>IIII')  # 定位到data开始的位置
    image_num = head[1]  # 获取图片数量
    width = head[2]
    height = head[3]
    bits = image_num * width * height 
    bits_string = '>' + str(bits) + 'B'  # fmt格式:'>47040000B'
    imgs = struct.unpack_from(bits_string, buffers, offset)  # 取data数据,返回一个元组
    binfile.close()
    imgs = np.reshape(imgs, [image_num, width * height])  # reshape为[60000,784]型的数组
    print ("load imgs finished")
    return imgs, head

其次,是加载标签,这是为了后续训练测试等工作。

# 加载标签
def load_label_set(filename):
    print ("load lable set",filename)
    binfile = open(filename, 'rb')  # 读取二进制文件
    buffers = binfile.read()
    head = struct.unpack_from('>II', buffers, 0)  # 读取label文件前两个整形数
    label_num = head[1]
    offset = struct.calcsize('>II')  # 定位到label数据开始的位置
    num_string = '>' + str(label_num) + 'B'  # fmt格式:'>60000B'
    labels = struct.unpack_from(num_string, buffers, offset)  # 取label数据
    binfile.close()
    labels = np.reshape(labels, [label_num])
    print ("load lable finished")
    return labels, head

这里还要进行 One-Hot 编码,常见处理了。

# 手动 one_hot 编码
def encode_one_hot(labels):
    num = labels.shape[0]
    res = np.zeros((num, 10))
    for i in range(num):
        res[i, labels[i]] = 1  # labels[i]表示0123456789,则 对应的列是1,这就是 One-Hot 编码
    return res

train_image = '/home/ma-user/work/Data/train-images.idx3-ubyte'
train_label = '/home/ma-user/work/Data/train-labels.idx1-ubyte'
test_image = '/home/ma-user/work/Data/t10k-images.idx3-ubyte'
test_label ='/home/ma-user/work/Data/t10k-labels.idx1-ubyte'
imgs, data_head = load_image_set(train_image)

# 这里的 label 是 60000 个数字,需要转成 one-hot 编码
labels, labels_head = load_label_set(train_label)
test_images, test_images_head = load_image_set(test_image)
test_labels, test_labels_head = load_label_set(test_label)

image.png

模型搭建/计算Loss/梯度更新

这一部分就是主要代码了,和在 GPU 上没什么区别,此部分代码一般无需改造。因此,迁移起来也不是很难。

# 定义参数
learning_rate = 0.01
training_epoches = 10
bacth_size = 100  # mini-batch
display_step = 2 # display once every 2 epochs

# tf graph input
x = tf.placeholder(tf.float32, [None, 784])  # 28 * 28 = 784
y = tf.placeholder(tf.float32, [None, 10])  # 0-9 ==> 10 classes

# 定义模型参数
W = tf.Variable(tf.zeros([784, 10]))  # tf.truncated_normal()
b = tf.Variable(tf.zeros([10]))

# 构建模型
prediction = tf.nn.softmax(tf.matmul(x, W) + b)
loss = tf.reduce_mean(-tf.reduce_sum(y * tf.log(tf.clip_by_value(prediction,1e-8,1.0)), reduction_indices=1))
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)
init = tf.global_variables_initializer()
res = encode_one_hot(labels)
print("res", res)
total_batchs = int(data_head[1] / bacth_size)
print("total_batchs:", total_batchs)

image.png

创建 session ,迁移到 NPU 并执行训练

注意,这里是迁移的关键了,简单来说,就是重写 config,将 device 设置为 NPU,就是将网络放到 NPU 上去执行,我们需要在创建 session 前添加如下配置,创建 config 并添加 custom_op:

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF

创建好的 config 作为 session config 传给 tf.Session ,使得训练能够在 NPU 上执行,sess.run 代码无需修改。

#训练
def train():
    with tf.Session(config=config) as sess:
            sess.run(init)
            for epoch in range(training_epoches):
                start_time = time.time()
                avg_loss = 0.
                total_batchs = int(data_head[1] / bacth_size)  # data_head[1]是图片数量

                for i in range(total_batchs):
                    batch_xs = imgs[i * bacth_size: (i + 1) * bacth_size, 0:784]
                    batch_ys = res[i * bacth_size: (i + 1) * bacth_size, 0:10]

                    _, l = sess.run([optimizer, loss], feed_dict={x: batch_xs, y: batch_ys})

                    # 计算平均损失
                    avg_loss += l / total_batchs
                end_time = time.time()
                if epoch % display_step == 0:
                    print("Epoch:", '%04d' % (epoch), "loss=", "{:.9f}".format(avg_loss), "time=", "{:.3f}".format(end_time-start_time) )

            print("Optimization Done!")

            correct_prediction = tf.equal(tf.argmax(prediction, 1), tf.argmax(y, 1))
            accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

            print("Accuracy:", accuracy.eval({x: test_images, y: encode_one_hot(test_labels)}))
train()

image.png

现在,这个 MNIST 模型的准确度已经达到约 74.48%。 在创建 session 时,除了以上基本配置外,其实,还可以通过相关配置,在 NPU 上使能混合计算、Profiling性能数据采集、训练迭代训练下沉等能力,想要了解更多,可以登录昇腾社区(https://www.hiascend.com) 阅读相关文档。

我们可以在 Terminal 执行如下命令查看 NPU 的情况:

watch -n 0.1 npu-smi info

这个操作类似于 GPU 的:

watch -n 0.1 nvidia-smi

看看这两条命令是不是很像?运行效果也比较类似,NPU 的运行效果如下图所示:

image.png

可以看到 AI Core 利用率,这个是我们主要关心的,因为这是 NPU 的算力核心,一般利用率越高,表示 NPU 的硬件性能发挥越好,一般也就越快了。值得一提的是,这里使用的是 Ascend 910A,高达32GB HBM,这可是高端显卡采用的显存啊,我记得优势是带宽大,好像速度慢了些,但一般 AI 训练的瓶颈是带宽,速度慢点可以接受。当然,我们也可以通过 ModelArts 右侧的资源监控来查看,可视化效果很好,赞一下。

image.png

使能混合精度

先来谈谈为什么要使用混合精度。简单来说,就是在不牺牲精度的情况下,为了节省内存(也可以说节省显存),加快训练速度。

混合精度训练方法是通过混合使用 float16 和 float32 数据类型来加速深度神经网络训练的过程,并减少内存使用和存取,从而可以训练更大的神经网络。同时又能基本保持使用 float32 训练所能达到的网络精度。当前昇腾 AI 处理器支持多种训练精度模式,用户可以在训练脚本中设置。不过,需要说明的是,个人感觉由于Ascend 910 的达芬奇架构以及整体设计,对 float32 的支持不太友好,甚至有些算子不支持 float32,这些在文档中有说明,可以详细看看,也期待下一代产品能够有所优化和改进。

下面来具体说说如何使能混合精度,得益于华为工程师们的努力,我们的操作变得很简单,只需要在前面的 config 中设置"precision_mode"为"allow_mix_precision"。

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
custom_op.parameter_map["precision_mode"].s = tf.compat.as_bytes("allow_mix_precision")  # 使能混合精度

再次执行训练部分,此时已经使能了混合精度.

# 训练
train()

image.png

是不是很简单,一行代码就搞定了!一般来说,混合精度可以提升训练速度,相应地内存和“显存”占用也会减少。不过这里效果不太明显,应该是 网络太小,即使不做优化,Ascend 910 强大的性能也足够支撑了。

使能 Loss Scaling

在混合精度计算中使用 float16 数据格式会导致数据动态范围降低,造成梯度计算出现浮点溢出,会导致部分参数更新失败。为了保证部分模型训练在混合精度训练过程中收敛,需要配置 Loss Scaling 的方法。

Loss Scaling 方法通过在前向计算所得的 loss 乘以 loss scale 系数 S,起到在反向梯度计算过程中达到放大梯度的作用,从而最大程度规避浮点计算中较小梯度值无法用 float16 表达而出现的溢出问题。在参数梯度聚合之后以及优化器更新参数之前,将聚合后的参数梯度值除以 loss scale 系数 S 还原。
Loss Scaling 分动态和静态两种方法:

动态 Loss Scaling

通过在训练过程中检查梯度中浮点计算异常状态,自动动态选取 loss scale 系数 S 以适应训练过程中梯度变化,从而解决人工选取 loss scale 系数 S 和训练过程中自适应调整的问题。

静态 Loss Scaling

顾名思义就是 Loss Scaling 初始化后保持不变,需要开发者能够人工选择一个适合网络的值。

这里主要和大家一起实现一个简单的动态 Loss Scaling。
相对之前的迁移,需要额外创建一个 NPULossScaleOptimizer,并实例化一个 ExponentialUpdateLossScaleManager 类进行动态 Loss Scale 的配置。

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF
custom_op.parameter_map["precision_mode"].s = tf.compat.as_bytes("allow_mix_precision") #使能混合精度

optimizer = tf.train.GradientDescentOptimizer(learning_rate) # NPULossScaleOptimizer接受tf的optimizer,这里重新初始化下optimizer
loss_scale_manager = ExponentialUpdateLossScaleManager(init_loss_scale=2**32, \
                                                       incr_every_n_steps=1000, decr_every_n_nan_or_inf=2, decr_ratio=0.5)#实例化1个loss manager
optimizer = NPULossScaleOptimizer(optimizer, loss_scale_manager) #创建loss scale optimizer
optimizer = optimizer.minimize(loss) #传入sess 执行的最后是loss的minimize值

再次执行训练

train()

image.png

令人欣喜的是,精度还有小幅度提升!

好了,今天就介绍到这里,更多详情,可参考https://www.hiascend.com/。

【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区),文章链接,文章作者等基本信息,否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件至:cloudbbs@huaweicloud.com进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容。
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。