How to measure the training cost(FLOPs)?

[feiga]

When I make some change about the model architecture, I would like to know how the training cost changes.
So I'm wondering is there any code in Tensor2Tensor (or any API of tensorflow) can measure the training cost (FLOPs), as shown in Table 2 of paper "Attention is All You Need" ? Besides, is it possible to profile the running time of each layer in the network architecture?

Thanks a lot.

[wenkesj]

@feiga closest/easiest thing I can think of, if you can't find it here, is tfprof.

[lukaszkaiser]

There is a timeline feature in TensorFlow. You need to generate it in the right format and then you can see beautiful graphs in Chrome, I do that quite a bit. See here:
https://stackoverflow.com/questions/34293714/can-i-measure-the-execution-time-of-individual-operations-with-tensorflow

The question would be how/where exactly to add it in tensor2tensor. We don't want to run profiling in every step (it's slow) but also not too rarely. Maybe in trainer_utils? Help would be welcome!

[feiga]

@lukaszkaiser Thanks!

Yes, I know the way you mentioned. It seems that we have to generate a run_metadata.

The tensor2tensor code uses tf.contrib.learn.learn_runner and tf.contrib.learn.Experiment to trigger the training of TF graph. I only find the session.run API to get the run_metadata. I don't know how to get run_metadata with tf.contrib.learn package. Do you have any ideas? I'd like to have a try.

[lukaszkaiser]

This is a very good question! I think you should open an issue against TensorFlow -- there should be a way of getting Timeline with tf.Estimator, but I don't know how to do it without a special support at this point. I'll ask around, but opening an issue would certainly be good.

[lukaszkaiser]

Ok, there is a session hook that does sth like that in a test here:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/monitored_session_test.py#L785

Could you try to copy-paste this hook code in trainer_utils and add it to the hooks in Estimator and check if it works?

[feiga]

Thanks, Lukasz! It works.
Here is a graph I get. Very cool feature of TensorFlow.

[lukaszkaiser]

Cool :). How about making a PR with a flag to allow others to see it too? Thanks!