TensorFlow: Как измерить, сколько памяти GPU занимает каждый тензор?

В настоящее время я внедряю YOLO в TensorFlow, и я немного удивлен тем, сколько памяти занимает. На моем GPU я могу обучать YOLO, используя их среду Darknet с размером пакета 64. На TensorFlow я могу делать это только с размером пакета 6, с 8 у меня уже не хватает памяти. Для фазы тестирования я могу работать с размером пакета 64 без исчерпания памяти.

  1. Мне интересно, как я могу рассчитать, сколько памяти потребляется каждым тензором? Все ли тензоры по умолчанию сохранены в графическом процессоре? Могу ли я просто рассчитать общее потребление памяти как форму * 32 бита?

  2. Я заметил, что, поскольку я использую импульс, все мои тензоры также имеют /Momentum тензор. Может ли это также использовать много памяти?

  3. Я дополняю свой набор данных методом distorted_inputsочень похоже на то, что определено в учебном пособии CIFAR-10. Может быть, эта часть занимает огромный кусок памяти? Я считаю, что Darknet делает изменения в процессоре.

3 ответа

Решение

Теперь, когда 1258 закрыто, вы можете включить регистрацию памяти в Python, установив переменную среды перед импортом TensorFlow:

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

В результате этого будет много регистрации. Вы хотите просмотреть результаты, чтобы найти соответствующие строки. Например:

grep MemoryLogTensorAllocation train.log

Извините за медленный ответ. К сожалению, сейчас единственный способ установить уровень журнала - это изменить файл tenorflow/core/platform/logging.h и перекомпилировать его, например:

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

Существует открытая ошибка 1258 для более элегантного управления журналом.

MemoryLogTensorOutput Записи регистрируются в конце каждого исполнения Op и указывают тензоры, которые содержат выходы Op. Полезно знать эти тензоры, поскольку память не освобождается до тех пор, пока нижестоящая операционная система не использует тензоры, что может быть намного позже на большом графике.

Смотрите описание в этом ( коммит). Распределение памяти - это необработанная информация, хотя для сбора информации в удобной для чтения форме необходим сценарий.

Другие вопросы по тегам