Блог

TensorFlow всегда (предварительно) выделяет всю свободную память (VRAM) на моей видеокарте, что нормально, поскольку я хочу, чтобы мои симуляции выполнялись как можно быстрее на моей рабочей станции.

Тем не менее, я хотел бы зарегистрировать, сколько памяти (в сумме) TensorFlow действительно использует. Кроме того, было бы действительно неплохо, если бы я мог также регистрировать, сколько памяти используют отдельные тензоры.

Эта информация важна для измерения и сравнения объема памяти, необходимого различным архитектурам ML / AI.

Есть какие-нибудь советы?

Обновление, можно использовать TensorFlow ops для запроса распределителя:

 # maximum across all sessions and .run calls so far
sess.run(tf.contrib.memory_stats.MaxBytesInUse())
# current usage
sess.run(tf.contrib.memory_stats.BytesInUse())
  

Также вы можете получить подробную информацию о session.run вызове, включая всю память, выделяемую во время run вызова, просмотрев RunMetadata . Т.е. что-то вроде этого

 run_metadata = tf.RunMetadata()
sess.run(c, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True), run_metadata=run_metadata)
  

Вот сквозной пример — возьмите вектор столбца, вектор строки и добавьте их, чтобы получить матрицу дополнений:

 import tensorflow as tf

no_opt = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0,
                             do_common_subexpression_elimination=False,
                             do_function_inlining=False,
                             do_constant_folding=False)
config = tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=no_opt),
                        log_device_placement=True, allow_soft_placement=False,
                        device_count={"CPU": 3},
                        inter_op_parallelism_threads=3,
                        intra_op_parallelism_threads=1)
sess = tf.Session(config=config)

with tf.device("cpu:0"):
    a = tf.ones((13, 1))
with tf.device("cpu:1"):
    b = tf.ones((1, 13))
with tf.device("cpu:2"):
    c = a b

sess = tf.Session(config=config)
run_metadata = tf.RunMetadata()
sess.run(c, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True), run_metadata=run_metadata)
with open("/tmp/run2.txt", "w") as out:
  out.write(str(run_metadata))
  

Если вы откроете run.txt , вы увидите сообщения, подобные этому:

   node_name: "ones"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322108320
      }
  ....

  node_name: "ones_1"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322092992
      }
  ...
  node_name: "add"
      allocation_description {
        requested_bytes: 676
        allocator_name: "cpu"
        ptr: 4492163840
  

Итак, здесь вы можете видеть, что a и b выделили по 52 байта каждому (13 * 4), а в результате выделили 676 байт.

Ответ Ярослава Булатова — лучшее решение для TF1.

Однако для TF2 contrib пакет не существует. Лучший способ — использовать профилировщик tf — https://www.tensorflow.org/guide/profiler#memory_profile_tool

Он построит график использования памяти, подобный этому.