【他山之石】tensorflow2.4性能调优最佳实践

“他山之石，可以攻玉”，站在巨人的肩膀才能看得更高，走得更远。在科研的道路上，更需借助东风才能更快前行。为此，我们特别搜集整理了一些实用的代码链接，数据集，软件，编程技巧等，开辟“他山之石”专栏，助你乘风破浪，一路奋勇向前，敬请关注。

作者：知乎—star all

地址：https://www.zhihu.com/people/star-all

01

开始

最近想尝试一下用Transformer做图片分类的效果，于是就在网上找找有没有比较好的例子．发现keras官方有个例子，于是就clone下来看看．本以为multi-head-attention这个模块需要自己来实现，竟然发现tf.keras中已经实现了multi-head-attention的接口，发现是真的方便（tensorflow的最新版本tf2.4才有的一个接口）．

跑了一个官方给的cifar的例子，效果还行．于是就打算在自己的数据上跑跑看效果，在这个过程中，发现官方给你例子还远远达不到训练速度最优化的程度．于是就把官方例子就改了一下，最终达到了一个满意的训练速度，本篇就是记录一下tf2.4下的训练性能调优全过程．最终调优后的代码会共享出来．

02

单卡到多卡

官方给的vit(Vision Transformer)的例子是基于单卡的，但是现在训练大型网络已经离不开多卡，于是把官方例子改成多卡训练的版本（这里主要说明一下单机多卡的设置）．

首先根据你手上卡的数量来建立一个strategy的对象．

physical_devices = tf.config.list_physical_devices('GPU')          
for ind, item in enumerate(physical_devices):                      
    tf.config.experimental.set_memory_growth(item, True)           

TRAIN_GPUS = [0,1,2,3]                                          
devices = ["/gpu:{}".format(i) for i in TRAIN_GPUS]             
strategy = tf.distribute.MirroredStrategy(devices)

这里tf.config.experimental.set_memory_growth的作用是限制显存的使用．

然后我们要把数据做一个副本化的包装：

train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)   
test_dist_dataset = strategy.experimental_distribute_dataset(test_daset)

然后我们需要对train_step和test_step做一个封装：

with strategy.scope()
        def distributed_train_step(dataset_inputs):                           
            per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))                           
            return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)               
        def distributed_test_step(dataset_inputs):                           
            return strategy.run(test_step, args=(dataset_inputs,))

最后在train_step，test_step，和计算loss和accuracy加上strategy的scope（这里只是拿train_step和test_step举例）：

with strategy.scope():                                                       
    def train_step(inputs):                                                  
        images, labels = inputs                                              
        with tf.GradientTape() as tape:                                      
            predictions = model(images, training=True)                       
            loss = compute_loss(labels, predictions)                         

        gradients = tape.gradient(loss, model.trainable_variables)           
        optimizer.apply_gradients(zip(gradients, model.trainable_variables)) 
        compute_acc(labels, predictions, train_accuracy)                     
        return loss                                                          

    def test_step(inputs):                                                   
        images, labels = inputs                                              
        predictions = model(images, training=False)                          
        compute_acc(labels, predictions, test_accuracy)

经过了上述的步骤，成功的在tf2.4下，把单卡训练转换到了单机多卡训练．

03

eager mode 到 static graph mode

由于tf2.x模式的执行方式是急切执行（eager mode），eager mode的好处在在于方便debug，但是如果拿来训练就不太好了，因为eager mode会托慢速度．所以我们需要在调试好网络之后把执行模式切换为静态图的模式．而这个步骤非常简单，加上一个tf.function的修饰符就好．

with strategy.scope():                                                   
        @tf.function                                                         
        def distributed_train_step(dataset_inputs):                           
            per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))         
            return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)             
        @tf.function                                                         
        def distributed_test_step(dataset_inputs):                           
            return strategy.run(test_step, args=(dataset_inputs,))

04

调整tf.data数据流顺序

tf.data的数据处理顺序非常重要，改变顺序可能会极度的托慢训练速度．

这里，我的数据处理流程如下：

1. 读取所有图片的路径和对应的label．
2. 把图片路径给parse成图片．

流程其实很简单，但这里要涉及到几点：

1. 如何shuffle.
2. 哪里设定epoch
3. 哪里设定batch
4. 哪里设定预取prefetch．

先来上code，再解释：

image_roots, labels = generate_fileroots_labels(file_root)                   
    dataset = tf.data.Dataset.from_tensor_slices((image_roots, labels))           
    dataset = dataset.repeat(100).shuffle(buffer_size=2000) 
  #  dataset = dataset.map(_parse_data, num_parallel_calls=tf.data.experimental.AUTOTUNE) 
    dataset = dataset.map(_parse_data, num_parallel_calls=16)                     
    dataset = dataset.batch(batch_size)                                           
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)

首先，repeat表示需要把数据重复多少次，也就是设定的epoc，shuffle代表在多少的buffer中打乱数据．这两者需要放到map前面，因为在map前，数据流处理的都是图片路径和label的轻量化数据，这对于repeat和shuffle是有利的．

而batch和prefetch就需要放到map的后面．这里需要注意，要先设定batch，再prefetch（不然会慢）．

（感谢评论的提醒，这里需要说明一下，官方文档里面是建议把map放到batch的前面，主要是因为map处理一个batch比处理单一的图片的效率高，理论上速度也会更快．但是我这边的map是处理单张图片，这种情况下就必须要把map放到batch的前面．）

05

大杀器－－tf.profiler

当我们觉得已经把加速做到极致了之后，我们需要用tensorflow自带的性能检测工具tf.profiler来检查一些性能还有那些可以榨取的空间．

用tf.profiler之前，你需要先按照官方的教程安装．这里有个小坑，因为tf.profiler需要依赖libcupti这个库，而libcupti这个库不在cuda的主库目录里，而是在extras/CUPTI/lib64里面，这个需要注意．

然后在你的训练代码中，需要添加如下的code：

for t_step, x in enumerate(train_dist_dataset):
        if t_step == 500:
              tf.profiler.experimental.start('/tmp/' + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
        if t_step == 600:
              tf.profiler.experimental.stop()
        with tf.profiler.experimental.Trace('Train', step_num=t_step, _r=1):
            step_loss = distributed_train_step(x)

这个代码段的意思是，在训练的第500步到第600步需要记录你的训练profile．

这里说明两点：

第一，我们不需要要整个训练过程都记录profile，因为记录profile仅仅是为了调优，只需要记录某些步的profile就可以提供你来调优即可．

第二，不从第０步就开始记录是因为我们需要先让训练达到稳定之后记录才会比较准确（示例中是从500步到600步开始记录profile）．

好了，我们看一下，都记录了一些什么东西：

这里我们可以看到耗时主要在3个方面：

1. Kernel Launch Time
2. Host Compute Time
3. Device Compute Time

我们来一个一个解决，首先来解决kernel lanuch time，这个在右边的建议(Recommendation for Next Step)有说明：

14.3 % of the total step time sampled is spent on 'Kernel Launch'. It could be due to CPU contention with tf.data. In this case, you may try to set the environment variable TF_GPU_THREAD_MODE=gpu_private.

也就是说可以通过设定TF_GPU_THREAD_MODE=gpu_private来解决．也就是说要在之前训练程序的前面加上下面一句命令：

export TF_GPU_THREAD_MODE=gpu_private

然而我这样执行之后发现作用好像不是很大，知道怎么解决的同学可以交流．

我们现在来解决第二个主要耗时的点，也就是Host Compute Time耗时过高．

我们知道，tensorflow的多卡策略是ps-worker的方式，这个ps可以认为是host，主要是负责更新参数和处理数据流，按理说这部分的耗时不应该很高才对．于是我往下看具体的host的耗时的页面：

我们看到，host的耗时，很大的程度和２个op有关：

1. stridedSlice
2. cast

其中stridedSlice占了绝对的大头，经过查阅资料发现，stridedSlice耗时比较高主要和tf.distribute.MirroredStrategy这个对象有关．

因为我在我的工程里面用到了tf.data.Dataset.from_tensor_slices这个对象，这个对象用到了stridedSlide这个操作，而tf.distribute.MirroredStrategy对stridedSlice的操作支持的不好．

网上的建议是把tf.distribute.MirroredStrategy换成tf.distribute.experimental.MultiWorkerMirroredStrategy．

也就是如下实现：

#  TRAIN_GPUS = [0,1,2,3]
#  devices = ["/gpu:{}".format(i) for i in TRAIN_GPUS]
#  strategy = tf.distribute.MirroredStrategy(devices)
tf.config.set_visible_devices(physical_devices[0:8], 'GPU') 
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

更换了之后，果然Host Compute Time 得到了下降．

关于cast这个操作，是因为我在dataset的parse_function里面用到了tf.cast的操作，我把这个操作放到了网络里面，这部分的耗时也消除了（其实是分给worker了）．

关于Device Compute Time这一部分，我发现这个部分主要的耗时用到了矩阵操作Einsum上，这一部分的操作也是MultiHeadAttention的主要操作，于是也就没有修改这个部分．

06

最后

当然，说了这么多，不如大家来下面一行代码来的方便.....

import torch as tf

不过话说回来，tensorflow和pytorch不是二选一的问题，而是我全都要，这样才能更好在深度学习里探(lian)索(dan)．

最后把示例代码放到这里，大家自行取用．

代码地址：https://github.com/Zhengtq/tensorflow2.x-best-practice

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“他山之石”历史文章

更多他山之石专栏文章，

请点击文章底部“阅读原文”查看

分享、点赞、在看，给个三连击呗！