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

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

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

目录

一、工具

二、结果展示

三、Show Code

四、经验


一、工具

timeline

二、结果展示

几行代码分析TensorFlow训练模型耗时_第1张图片

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

几行代码分析TensorFlow训练模型耗时_第2张图片

图片来源:使用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 文件:

几行代码分析TensorFlow训练模型耗时_第3张图片

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

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

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

以上,问题解决。

你可能感兴趣的:(三妹的工作日常,深度学习,机器学习,tensorflow,模型)