Tensorflow Operation에 대한 프로파일링 방법

tips
tensorflow

(Curtis Kim) #1

session.run() 시점에 options 파라미터로 timeline을 추가하여 저장하면 실제 오퍼레이션에 걸린 시간을 타임라인으로 Visualization 할 수 있습니다.

import tensorflow as tf
from tensorflow.python.client import timeline

x = tf.random_normal([1000, 1000])
y = tf.random_normal([1000, 1000])
res = tf.matmul(x, y)

# Run the graph with full trace option
with tf.Session() as sess:
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    sess.run(res, options=run_options, run_metadata=run_metadata)

    # Create the Timeline object, and write it to a json
    tl = timeline.Timeline(run_metadata.step_stats)
    ctf = tl.generate_chrome_trace_format()
    with open('timeline.json', 'w') as f:
        f.write(ctf)

위와 같이 sess.run() 시점에 run_options를 넣어주고, step_stats를 읽어서 json으로 저장하는 형태입니다.

이렇게 저장된 파일은 브라우저인 크롬에서 지원하는 형태이므로, 크롬에서 chrome://tracing 주소로 접속 한 후 해당 파일을 읽으면 아래와 같이 타임라인을 프로파일링할 수 있습니다.

https://github.com/ildoonet/tf-lcnn/raw/master/images/timeline_alexnet.png