TensorFlow: How to measure how much GPU memory each tensor takes?

Question

I\'m currently implementing YOLO in TensorFlow and I\'m a little surprised on how much memory that is taking. On my GPU I can train YOLO using their Darknet framework with b

Answer 1

See the description in this (commit). The memory allocation is raw info is there although it needs a script to collect the information in an easy to read form.

Answer 2

Sorry for the slow reply. Unfortunately right now the only way to set the log level is to edit tensorflow/core/platform/logging.h and recompile with e.g.

#define VLOG_IS_ON(lvl) ((lvl) <= 1)

There is a bug open 1258 to control logging more elegantly.

MemoryLogTensorOutput entries are logged at the end of each Op execution, and indicate the tensors that hold the outputs of the Op. It's useful to know these tensors since the memory is not released until the downstream Op consumes the tensors, which may be much later on in a large graph.

Answer 3

Now that 1258 has been closed, you can enable memory logging in Python by setting an environment variable before importing TensorFlow:

import os
os.environ['TF_CPP_MIN_VLOG_LEVEL']='3'
import tensorflow as tf

There will be a lot of logging as a result of this. You'll want to grep the results to find the appropriate lines. For example:

grep MemoryLogTensorAllocation train.log