Bash Python Profiling for TensorBoard

.bashrc.tensorboard

# -*-sh-*-

# pip install yappi
# tensorboard-profile-wall-time --logdir=/tmp/mnist
tb-profile-wall-time() {
  -tensorboard-profile wall "$@"
}

tb-profile-cpu-time() {
  -tensorboard-profile cpu "$@"
}

-tb-profile() {
  clock=$1
  shift
  bazel build //tensorboard || return
  # yappi -o profile.txt bazel-bin/tensorboard/tensorboard "$@"
  python <(cat <<EOF
import sys
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_html5lib')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_mozilla_bleach')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_pocoo_werkzeug')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_pythonhosted_markdown')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_pythonhosted_six')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/org_tensorflow_tensorboard')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/protobuf')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles/protobuf/python')
sys.path.insert(0, 'bazel-bin/tensorboard/tensorboard.runfiles')
import yappi
import atexit
@atexit.register
def done():
  yappi.convert2pstats(yappi.get_func_stats()).dump_stats('profile.binary')
yappi.set_clock_type('$clock')
yappi.start()
from tensorboard import main
main.main()
EOF
) "$@"
  python <(cat <<'EOF'
import os
import pstats
with open('profile.txt', 'w') as out:
  pstats.Stats('profile.binary', stream=out).print_stats()
os.unlink('profile.binary')
EOF
)
  sed -i -e 's! [^ ]*/tensorboard.runfiles/! !' profile.txt
  sed -i -e 's! [^ ]*/site-packages/! !' profile.txt
  (
    sed -n 1,7p profile.txt
    head -n20 profile-ncalls.txt
    echo '... [see profile-ncalls.txt]'
    echo
    echo '=== TOTAL TIME ==='
    sed -n 7p profile.txt
    sed 1,7d profile.txt | sort -nrk2 >profile-tottime.txt
    head -n20 profile-tottime.txt
    echo '... [see profile-tottime.txt]'
    echo
    echo '=== TOTAL TIME PER CALL ==='
    sed -n 7p profile.txt
    sed 1,7d profile.txt | sort -nrk3 >profile-tottime-percall.txt
    head -n20 profile-tottime-percall.txt
    echo '... [see profile-tottime-percall.txt]'
    echo
    echo '=== CUMULATIVE TIME ==='
    sed -n 7p profile.txt
    sed 1,7d profile.txt | sort -nrk4 >profile-cumtime.txt
    head -n20 profile-cumtime.txt
    echo '... [see profile-cumtime.txt]'
    echo
    echo '=== CUMULATIVE TIME PER CALL ==='
    sed -n 7p profile.txt
    sed 1,7d profile.txt | sort -nrk5 >profile-cumtime-percall.txt
    head -n20 profile-cumtime-percall.txt
    echo '... [see profile-cumtime-percall.txt]'
  ) | less -S
}