几行代码分析TensorFlow训练模型耗时

三妹最近被领导质疑了，他说我模型训练时间太久：你看人家XXX公司，千万级的数量级几个小时就训练好了，你这个快40个小时了，必须得优化，你先看看训练模型时间都花在哪了吧。

经过一顿尝试，我要记录一下我觉得最简单高效的方法。

目录

一、工具

二、结果展示

三、Show Code

四、经验

---

一、工具

timeline

二、结果展示

 图片来源：tensorflow性能调优实践 - 简书

图片来源：使用TensorFlow训练WDL模型性能问题定位与调优 - 美团技术团队 

（P.S. 感谢大佬们的无私分享）

三、Show Code

import tensorflow as tf

# 以下为重要代码片段，一般写在 train() 函数中
with tf.Session() as sess:
    # 定义 run_options 和 run_metadata, 用于保存op的属性
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()

    from tensorflow.python.client import timeline
    
    # 训练 10W step
    for i in range(100000):
        sess.run(train_op, options=run_options, run_metadata=run_metadata)
        if i % 1000 == 0:
            train_loss = sess.run(train_loss)
            sess.run(train_auc_op)
            train_auc = sess.run(train_auc_value)
                
            print("Step:", i, "train_loss:", train_loss, "train_auc:", train_auc)
            
            # 记录耗时
            fetched_timeline = timeline.Timeline(run_metadata.step_stats)
            chrome_trace = fetched_timeline.generate_chrome_trace_format()
            with open('./timeline/timeline_train_{}.json'.format(i), 'w') as f:
                    f.write(chrome_trace)

四、经验

1、耗时文件写成 JSON格式，保存之后，浏览器中输入：chrome://tracing，点击 load 后选中保存好的 JSON 文件：

2、代码中任何 sess.run() 函数中都可以加记录耗时的参数配置，但是如果都用同一个run_options 和 run_metadata 的话，保存 JSON 时只会保存最后一个 sess.run() 的耗时（别问我怎么知道的，都是坑啊🥲）

3、我的代码是记录多个 step 的结果，不过按照我的经验看，第100个step的耗时和第100000个step的耗时分布几乎是一致的。

4、load 之后就可以看到所有过程的耗时，再去优化耗时比例最大的部分（我的是 DecodeCSV，等我优化成功、得到组长认可之后再来更文呀，你们会等我的对嘛mua）

以上，问题解决。