训练时添加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上述文件即可查看。