Использование памяти ЦП увеличивается с каждым пакетом, загруженным с использованием tf.data.Dataset и tf.py_function - PullRequest
Новая
       7

Использование памяти ЦП увеличивается с каждым пакетом, загруженным с использованием tf.data.Dataset и tf.py_function

0 голосов
/ 10 октября 2019

У меня есть данные в больших файлах в произвольном формате. В данный момент я создаю tf.data.Dataset с именем этих файлов, а затем вызываю tf.py_function, чтобы получить к ним доступ, необходимый для моего обучения. Функция py_function должна загрузить весь файл (почти 8 ГБ) в память, чтобы создать массив из нескольких МБ (1024x1024x4). Функция py_function возвращает только этот массив и соответствующую метку. Проблема в том, что каждый загружаемый образец увеличивает нагрузку на ЦП почти на 8 ГБ. Очень быстро на моем компьютере заканчивается ОЗУ и происходит сбой программы. Когда я запускаю свою программу вне VS Code, я получаю вдвое больше пакетов, чем когда использую отладчик, но все равно это максимум 13 пакетов. (У меня 32 ГБ ОЗУ процессора и 13 * 8> 32, так что похоже, что память иногда освобождается, но, может быть, недостаточно быстро?)

Я сохраняю batch_size и prefetch оба маленькими, так что тольконемногие из этих больших массивов должны быть в памяти одновременно. Я ожидал, что тензор потока освободит эту память, как только выйдет функция py_function, и она выйдет из области видимости. Я пытался побудить освободить память раньше, явно удалив переменную и вызвав сборщик мусора, но это не помогло.

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

import pickle
import gc
import tensorflow as tf

def tf_load_raw_data_and_labels(raw_data_files, label_files):
    [raw_data, labels] = tf.py_function(load_raw_data_and_labels, [raw_data_files, label_files], [tf.float32, tf.float32])

    raw_data.set_shape((1024, 1024, 4))
    labels.set_shape((1024))
    return raw_data, labels

def load_raw_data_and_labels(raw_data_file, label_file):
    #load 8GB big_datacube, extract what I need into raw_data
    del big_datacube
    gc.collect() #no noticeable difference
    return raw_data, labels

with open("train_tiles_list.pickle", "rb") as fid:
    raw_data_files, label_files = pickle.load(fid)

train_dataset = tf.data.Dataset.from_tensor_slices((raw_data_files, label_files))
train_dataset = train_dataset.shuffle(n_train)#.repeat()
train_dataset = train_dataset.map(tf_load_raw_data_and_labels)
train_dataset = train_dataset.batch(batch_size)
train_dataset = train_dataset.prefetch(1)

Я обычно обучаю модель ResNet50, используя функцию model.fit() из tf.keras.Model, но я также пыталсянастройка из краткого руководства tf2 , которое позволяет установить точку отладки после обучения каждого пакета. В точке отладки я проверил использование памяти активными переменными . Список очень короткий, и нет переменных более 600 КБ. В этот момент gc.collect() возвращает число от 10 до 20, даже после его запуска несколько раз, но я не слишком уверен, что это значит.

Может оказаться, что проще всего будет пройти через всебольшие файлы и сохраните меньшие массивы в свои файлы, прежде чем я начну любое обучение. Но сейчас я хотел бы понять, есть ли что-то фундаментальное, что не освобождает память. Это утечка памяти? Возможно, связано с tf.data.Datasets, py_functions или чем-то другим, специфичным для моей установки?

Редактировать: Я читал, что сборка мусора в python была обновлена ​​с помощью python3.4 . Из-за зависимости, связанной с моими пользовательскими данными, я использую python2.7. Может ли это быть частью проблемы?

Редактировать 2: Я обнаружил некоторые проблемы с github по поводу утечек памяти при использовании tenorflow. Предложенный обходной путь (tf.random.set_seed(1)) не работает для меня:
https://github.com/tensorflow/tensorflow/issues/31253
https://github.com/tensorflow/tensorflow/issues/19671

...