前言

首先，若是你如今已经很熟悉tf.data+estimator了，能够把文章x掉了╮(￣▽￣””)╭python

可是！若是如今仍是在进行session.run(…)的话！尤为是苦恼于GPU显存都塞满了利用率却上不去的童鞋，这篇文章或许能够给你打开新世界的大门噢(￣∇￣)编程

若是发现通过一系列改良后训练效率大大提升了，记得回来给小夕发小红包(￣∇￣)api

不过，这并非一篇怒贴一堆代码，言（三）简（言）意（两）赅（语）就结束的CSDN文风的文章。。。因此伸手党们也能够X掉了╮(￣▽￣””)╭session

缘起

很早很早以前，在小夕刚接触tensorflow和使用GPU加速计算的时候，就产生过一个疑惑。为何显卡的显存都快满了，GPU利用率还显示这么低呢？好浪费呀，可是又迫不得已。当时GPU利用率100%的状况基本是仅存于一块显卡塞四、5个不费显存的小任务的状况。多线程

大部分状况下写出来的代码train起来后是这样的：app

能够看到，虽然显卡的显存都塞满了，可是显卡功率（最左边那一栏，114W和69W）和利用率（最右边那一栏，35%和38%）却远远没有达到极限。大部分人的想法是，算了算了这不重要，我去作实验了再见【wei笑】函数式编程

然而！若是你在作大型实验，train一次跑几天呢？这个细节会极大的影响你的实验效率和DDL到来前的实验次数！想一下，彻底同样的model和设置，你的代码要train一周，然而隔壁老王只须要train三天╮(￣▽￣””)╭函数

路人甲：我有256张显卡
小夕：好了这篇文章你能够X掉了
那么，咱们有没有可能一直这样呢：

oop

是否是这功率和利用率看起来难以想象！不要怀疑这是PS的图！这只是小夕的平常截图！tricks用的好GPU利用率掉不下来99%，然鹅代码写的足够蠢，也能够上不去5%！post

那么问题来了，究竟是什么致使的这个差别呢？
不要急，咱们来放大一下那些gpu利用率只有30%几的代码在训练时的gpu利用率的变化状况（好像句子有点长

watch -n 0.1 nvidia-smi

ps:（可能掉帧太严重了看着不连贯╮(￣▽￣"")╭，建议在本身的机器上试一下，会直观的多～）
看！是否是一会儿就发现问题啦？能够看到，其实gpu利用率并非一直在比较低的水平，而是颇有规律的周期性的从0涨到接近100再跌到0，再从新涨到100再跌回0。若是同时开着打印日志的窗口，你就会发现这个周期刚好跟每一个训练step的时长一致！也就是说，在每一个step，其实有一些时间并无花在GPU里，那固然就是花在cpu里啦。

那在cpu里干什么的呢？固然就是load下一个batch、预处理这个batch以及在gpu上跑出结果后打印日志、后处理、写summary甚至保存模型等，这一系列的花销都要靠cpu去完成。回顾一下咱们常写的代码：

create_graph()
create_model_saver()
create_summary_writer()
create_session()
do_init()
for i in range(num_train_steps):
    load_batch(...)                # cpu
    preprocess(...)                # cpu
    feed_dict = {...}              # cpu
    fetch_list = [...]             # cpu
    buf = session.run(fetch_list, feed_dict)    # gpu
    postprocess(buf)               # cpu
    print(...)                     # cpu
    if i % x == 0:
        summary_writer.write(...)  # cpu
    if i % xx == 0:
        model_saver.save(...)      # cpu

看，尤为是preprocess(…)任务比较重的话就容易致使代码在cpu里也要跑好一段时间，gpu利用率天然就会上不去并且呈现周期性变化啦。

那么有没有什么办法下降cpu时间，提升gpu时间呢？
一个很自（愚）然（蠢）的想法就是把一切训练代码都用tf的api重写不就好啦，甚至最外层的那个for i in range(num_train_steps)其实均可以用tf.while_loop重写呀。嗯，小夕还真的这么尝试过，而后发现

TF api这特喵的都是些什么鬼！各类跟numpy和python内置函数重名却行为不一致是什么鬼！卧槽这个api少了个参数我该怎么办？python里就一行代码就能搞定的事情我为何写了几十行？？

因此除了函数式编程的大牛，小夕极力的不建议重蹈覆辙！尤为是咱们这些遇到汇编会哭，看到Lisp会崩溃的90后小仙女！

因此没办法把整个train loop都描述进计算图了？

别怕别怕，好在后来其实tensorflow已经封装了一个特别好（多）用（坑）的上层API来把整个train loop都能轻松的封装在计算图中，从而实现超级高的GPU利用率和训练效率！

Estimator

不用管它为啥叫Estimator，只须要知道，它把咱们刚才想作的事情基本都给封装好了就行。把刚才的那个经典的写法搬过来

1. create_model()
2. create_model_saver()
3. create_summary_writer()
4. create_session()
5. do_init()
6. for i in range(num_train_steps):
7.      load_batch(...)                # cpu
8.      preprocess(...)                # cpu
9.      feed_dict = {...}              # cpu
10.     fetch_list = [...]             # cpu
11.     buf = session.run(fetch_list, feed_dict)    # gpu
12.     postprocess(buf)               # cpu
13.     print(...)                     # cpu
14.     if i % x == 0:
15.         summary_writer.write(...)  # cpu
16.     if i % xx == 0:
17.         model_saver.save(...)      # cpu

1-5行在estimator中都封装好啦，你只须要把相关配置塞进estimator的RunConfig就能够啦～

7-9行也封装好啦，你只须要把数据集载入和预处理的相关代码的函数塞给estimator.train的input_fn～

第10行也封装好啦，你只须要把要fetch的loss、train_op丢进estimator的EstimatorSpec～

第11行也封装好啦，你只须要把描述模型计算图的函数塞给estimator的model_fn~

第12-13行不用操心细节了，global_step和loss自动完成了，剩下的丢给tf.Print和LoggingTensorHook吧～

第14-17行不用你写了，自动完成了

╮(╯▽╰)╭

通过这么一顿折腾，咱们发现GPU利用率大大提升啦～直逼80%甚至90%。那么还有没有能够压榨的空间呢？

其实这时仔细一分析就会发现虽然estimator把大部分的代码写进计算图里了，可是从数据的载入和预处理依然是在cpu里串行进行呀，并且好比一个batch有128个样本，那么estimaor内部在run每一个step的时候仍是要等着这128个样本串行的处理完才行。这显然就是最后的瓶颈啦！有没有办法消除掉呢？·固然有，那就是

tf.data

TF的dataset API能够说让人又爱又恨了，它确实看似提供了一种把整个预处理都搬进计算图进行并行化处理的途径，可是！若是你真的彻底用tensorflow API来作复杂的预处理的话，真的会让人疯掉的QAQ所以，这里在用tf.data以前，小夕极力的建议先把数据集尽量的transform成预处理后的样子，包括作分词、作截断、作word2id等，不过padding和input_mask能够留在TF里面作，毕竟都只须要一行。

那作完这些预处理后，数据该怎么存储会更方便后续的读取和处理呢？最最最建议的方式仍是使用tf.records来存储，磁盘、内存的存储和IO效率都会相比传统方式更快一些，x和y也不用分开了。固然这样的惟一的坏处就是不能直接打开看数据集╮(￣▽￣””)╭毕竟数据集被作成了二进制文件。

可是实在比较懒不想用tf.record的话，那么小夕极力建议把x和y分开存储，而且尽可能让tf.data在读取数据的时候作完上面的那些必要的预处理，以避开难用的字符串基础操做API而且减轻训练时的cpu和内存压力。

tf.data还有一个很大的好处就是能够很自然的支持以streaming的方式读取数据，这样在面对大数据集时就不会发生数据load完后发现显卡被占的尴尬事件了╮(￣▽￣””)╭

好像讲了这么久，仍是没讲怎么用tf.data加速QAQ，来来来进入正题啦。

想一想哈，没用tf.data的时候，咱们写出来的代码实际跑起来就是这个样子的：

这也是文章开头小夕解释的为何gpu利用率上不去而且周期性变化的重要缘由。那么咱们能够不能够消除idle，像下面这样让prepare和train的过程并行进行呢？

固然能够！那就是

prefetch

从prefetch的意思就能够理解，那就是预先获取下一个step要load的batch。使用tf.data里面的叫作prefetch的神奇api就能够轻松完成啦，这个api里的参数buffer_size就是讲的是额外的fetch多少份，好比buffer_size=1，而后咱们要prefetch的是batch的话，那么模型每次prepare完一个batch后，就会自动再额外的prepare一个batch，这样下一个train step到来的时候就能够直接从内存中取走这个事先prepare好的batch啦。（详情见后面）

等下，看上图的话，有木有发现，若是prepare一个batch耗时很短的话确实两全齐美，可是若是耗时比较久，尤为一会儿prefetch好几个batch的话，一旦prepare的用时超过了train一个step的用时，那么每一个train step的性能就会受限于prepare的效率啦。放大一下这个问题的话以下图所示

看，prepare用时过久反而会致使train完一个step后gpu空闲了（虽然其实下个step的batch可能已经prepare好了）

那么能不能确保prepare阶段的用时小于train阶段的用时呢？

parallel mapping

一个很简单的想法固然就是让样本并行处理啦～若是batch size是128，prefetch size=1，那么准备一个batch要串行的跑128*2=256次的预处理，可是若是咱们开4个线程去跑，是否是就看起来快多啦。幸运的是咱们也不用本身手撸多线程了，tf.data.Dataset在map（预处理）函数里有一个参数num_parallel_calls，给这个参数赋值就能够并行parse啦。如图，

这样的话只要prefetch的buffer_size和map的num_parrellel_calls取得合适，基本就能够实现不间断的train啦，也就是几乎达到100%的GPU利用率！

好啦，思想明白了，代码就容易理解啦。不使用tf.record，直接从预处理好的纯文本格式的数据集load数据时的典型过程以下

def build_input(..):
    x = tf.data.XXDataset(..)
    x = x.map(..., num_parallel_calls=N)        # parellel

    y = tf.data.XXDataset(..)
    y = y.map(..., num_parallel_calls=N)

    dataset = tf.data.Dataset.zip((x, y))
    dataset = dataset.repeat(num_epochs)    
    if is_train:
        dataset = dataset.shuffle(..)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(buffer_size=1)   # prefetch
    iterator = dataset.make_xx_iterator()
    return iterator.get_next()

固然，若是用上tf.record后，就不用分别从x和y俩文件中读数据啦，感兴趣的童鞋可自行去了解一下。

机器学习如何提升GPU利用率

前言

缘起

Estimator

tf.data

prefetch

parallel mapping