三妹最近被领导质疑了,他说我模型训练时间太久:你看人家XXX公司,千万级的数量级几个小时就训练好了,你这个快40个小时了,必须得优化,你先看看训练模型时间都花在哪了吧。
经过一顿尝试,我要记录一下我觉得最简单高效的方法。
目录
一、工具
二、结果展示
三、Show Code
四、经验
timeline
图片来源:tensorflow性能调优实践 - 简书
图片来源:使用TensorFlow训练WDL模型性能问题定位与调优 - 美团技术团队
(P.S. 感谢大佬们的无私分享)
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)
以上,问题解决。