Я пытаюсь реализовать функцию python для Tensorflow, которая может запрашивать статистику памяти у распределителей памяти вниз на уровне абстракции c ++. Есть ли примеры функциональности в TF, которые я мог бы использовать, чтобы пересечь питон и разделить с ++? Я смотрю на вещи, связанные с SWIG, которые, кажется, как код делает это, но это загадочно. Конкретный пример мог бы иметь здесь большое значение.
Я хочу переработать некоторые алгоритмы выделения памяти, чтобы потенциально улучшить производительность более крупных моделей. Таким образом, возможность запрашивать статистику памяти между итерациями модели имеет решающее значение. Первоначально я начал исследовать создание пользовательской операции с тензорным потоком для получения статистики, но после обсуждения этого вопроса с моей командой мы решили, что наилучшим путем будет реализация функции Python passthru, которая может запрашивать статистику, которая будет вызываться из чего-то вроде, скажем,Обратный вызов Keras.
Это поля, которые я хотел бы передать пользователю TensorFlow. Расположен в allocator.h.
struct AllocatorStats {
int64 num_allocs; // Number of allocations.
int64 bytes_in_use; // Number of bytes in use.
int64 peak_bytes_in_use; // The peak bytes in use.
int64 largest_alloc_size; // The largest single allocation seen.
// The upper limit of bytes of user allocatable device memory, if such a limit
// is known.
absl::optional<int64> bytes_limit;
// Stats for reserved memory usage.
int64 bytes_reserved; // Number of bytes reserved.
int64 peak_bytes_reserved; // The peak number of bytes reserved.
// The upper limit on the number bytes of reservable memory,
// if such a limit is known.
absl::optional<int64> bytes_reservable_limit;