训练时添加tf.train.ProfilerHook记录各个op的计算资源。
hook = tf.train.ProfilerHook(
save_steps=1000,
output_dir=os.path.join("./ckpt/", "tracing"),
show_dataflow=True,
show_memory=True)
hooks = [hook]
train_spec = tf.estimator.TrainSpec(
hooks=hooks,
input_fn=lambda: fe.train_input_fn(train_files, batch_size), max_steps=100000)
文件./ckpt/tracing/timeline-1000.json即是存储的各种op计算资源信息,可以在浏览器中输入chrome://tracing,load上述文件即可查看。