Керасу не хватает памяти при запуске моих игрушечных примеров

Question

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

import keras.layers 
import keras.models 
import numpy as np
from keras.optimizers import Adam
import random 
import os
import cv2
import tensorflow as tf

TrainingDirectory="/home/thijser/host/ImageConverter/trainingdata" 
def main():
    model = define_feedforward_network1()
    compileModel(model)
    model.summary()
    train(model)

def define_feedforward_network1():
    inp = keras.engine.input_layer.Input(shape=(None,None,3))
    init = keras.initializers.RandomNormal(stddev=0.02)
    layer1=keras.layers.Conv2D(64, (4,4), activation='relu', padding='same', kernel_initializer=init)(inp)
    layer2 = keras.layers.Conv2D(128, (4,4), activation='relu', padding='same', kernel_initializer=init)(layer1)
    outp = keras.layers.Conv2D(3, (4,4), activation='relu', padding='same', kernel_initializer=init)(layer2)

    return keras.Model(input=inp,output=outp)


def compileModel(model):
    adam=keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, amsgrad=False)
    model.compile(loss="mse", optimizer=adam,
    metrics=["mse"])

def train(model):
    model.fit_generator(mygenerator(),steps_per_epoch=100,epochs=16)


def mygenerator():

    batch_features=None
    batch_labels=None
    while True:
        target=random.choice(os.listdir(TrainingDirectory))
        batch_features=cv2.imread(TrainingDirectory+'/'+target+"/input.jpg", cv2.IMREAD_COLOR)
        batch_labels=cv2.imread(TrainingDirectory+'/'+target+"/labelcol.png", cv2.IMREAD_COLOR)
        yield np.array([batch_features]), np.array([batch_labels])

main()  

Какая простая нейронная сеть для (в данном случае) превращения цветаизображения в изображения в градациях серого (просто чтобы почувствовать это). Его можно запустить, поместив несколько папок в папку trainingDirectory и поместив 1 изображение с именем labelcol.png и другой файл input.jpg (такого же размера, как и первый) внутри каждой из этих папок.

Однако, когда я запускаю его, он выходит из ошибки довольно быстро (прямо в конце первых 16 изображений в тренинге). Если я включу report_tensor_allocations_upon_oom, вместо этого я получаю segfault, который не дает мне полезной информации, однако я могу видеть, что использование памяти увеличивается примерно на 100-300 МБ за раз, предполагая, что целые тензоры не могут быть освобождены.

Кто-нибудь видит, что я делаю не так?

Answer 1

Вы не можете загрузить все свои тренировочные образы в оперативную память. Вам нужно будет создать собственный DataLoader, который загружает изображения в пакетном режиме в вашем цикле обучения.

Эта статья имеет большое руководство, которое мне очень помогло.