Keras (2.0.8) con Tensorflow (1.3) backend toma toda la RAM disponible

Keras (2.0.8) con Tensorflow (1.3) backend toma toda la RAM disponible

Estoy usando la biblioteca keras con tensorflow backend y CUDA habilitado. Ver salida de versiones de paquetes PIP:

Keras (2.0.8)
tensorflow-gpu (1.3.0)
tensorflow-tensorboard (0.1.8)

Tengo el siguiente código que crea el modelo VGG16 y carga los pesos de ImageNet:

def create_vgg16_model(target_size: tuple, n_classes: int):
    base = VGG16(include_top=False,
                 input_shape=target_size,
                 weights='imagenet')

    x = base.output
    x = Flatten()(x)
    x = Dense(n_classes, activation='softmax', name='top')(x)

    model = Model(inputs=base.input, outputs=x)
    for layer in model.layers[:-1]:
        layer.trainable = False

    model.compile(optimizer='adam', loss='categorical_crossentropy')
    return model

El entrenamiento del modelo va bien y nvidia-smi muestra que la memoria GPU se utiliza según sea necesario. Pero luego revisé la salida del comando top y esto es lo que veo:

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND                                                                 
 1268 ck        20   0  166288  31964  12416 S  29.5  0.1  13:05.39 Xtightvnc                                                               
32235 ck        30  10   32252   3700   3356 S   5.3  0.0   0:36.93 cwaves 
------------------------------------------------------------------------------    
32212 ck        20   0 27.485g 1.184g 190404 S   2.3  3.8   0:35.44 python  
------------------------------------------------------------------------------                                                                
26015 root      20   0       0      0      0 S   0.3  0.0   0:00.30 kworker/3:1                                                             
31754 ck        20   0   43168   3904   3080 R   0.3  0.0   0:04.45 top                                                                     
    1 root      20   0  185644   6204   3984 S   0.0  0.0   0:10.44 systemd                                                                 

Revisé el código con el depurador y me di cuenta de que la memoria se asigna en la siguiente función tomada de keras.backend.tensorflow_backed que crea un objeto tf.Session:

def get_session():        
    global _SESSION
    if tf.get_default_session() is not None:
        session = tf.get_default_session()
    else:
        if _SESSION is None:
            if not os.environ.get('OMP_NUM_THREADS'):
                config = tf.ConfigProto(allow_soft_placement=True)
            else:
                num_thread = int(os.environ.get('OMP_NUM_THREADS'))
                config = tf.ConfigProto(intra_op_parallelism_threads=num_thread,
                                        allow_soft_placement=True)
                # next line allocates ~28GB of RAM
                _SESSION = tf.Session(config=config)
        session = _SESSION
    if not _MANUAL_VAR_INIT:
        with session.graph.as_default():
            _initialize_variables()
    return session

Y esto sucede para todos los modelos disponibles, porque la memoria se asigna cuando se crea una sesión, antes de que comience el entrenamiento o la inicialización de las variables.

Sé que TF asigna toda la memoria GPU disponible (a menos que anule ConfigProto y/o ajuste sus variables de entorno), pero ¿hace lo mismo con la RAM? Es decir. parece que el marco está asignando toda la RAM que tengo en mi máquina, excepto una que ya está asignada por otros procesos.

¿Alguien detectó tal comportamiento con diferentes versiones de tensorflow o keras? ¿Crees que hay alguna manera de limitar de alguna manera la cantidad de memoria utilizada?


Actualización 1

Hace algún tiempo, el núcleo eliminó uno de mis scripts de entrenamiento con un error de falta de memoria después de 50-60 períodos de entrenamiento. Aunque la estadística de uso de memoria de GPU volátil muestra que también se usa. (No solo asignado, según tengo entendido).


Actualización 2

De acuerdo, la memoria virtual no es una métrica válida. Pero descubrí que el consumo de memoria aumenta casi linealmente durante el proceso de entrenamiento del modelo. Tengo el siguiente ciclo de entrenamiento:

def train_model(model, x, y):
    loss = model.train_on_batch(x, y)
    return loss


def train_model_42(model, x, y):
    # dummy function
    return 42.0


def training_loop():
    # training parameters
    target_size = 224, 224, 3
    batch_size = 128

    # generator yielding batches of file paths
    files_stream = FilesStream(folder=TRAIN_IMAGES, batch_size=batch_size)
    files_source = files_stream()

    # list of generators loading images from persistent storage
    gens = [
        image_loader(),
        augment_images(horizontal_flip=True),
        shuffle_samples(),
        normalize_images(target_size=target_size)
    ]

    # Model from keras.applications with replaced top layer
    model = get_model('resnet50').build(n_classes=n_classes)

    for epoch in range(1, 1001):
        epoch_loss = []
        for _ in range(files_stream.steps_per_epoch):
            for gen in gens:
                gen.send(None)
            processed = next(files_source)
            for gen in gens:
                processed = gen.send(processed)
            x, y = processed
            loss = train_model_42(model, x, y) # <-- this shows pic. 1
            # loss = train_model(model, x, y)    <-- this shows pic. 2                  
            epoch_loss.append(loss)
        avg_loss = sum(epoch_loss) / len(epoch_loss)
        print('Epoch %03d: train loss = %2.4f' % (epoch, avg_loss))

Cuando uso la función de entrenamiento ficticia, el gráfico de consumo de memoria parece que se muestra en pic 1: ingrese la descripción de la imagen aquí

Pero mientras se ejecuta un proceso de entrenamiento real, parece pic 2: ingrese la descripción de la imagen aquí

¿Por qué aumenta el consumo de memoria durante el proceso de entrenamiento? ¿Se almacenan en caché lotes de datos anteriores o algo así? ¿Debería el modelo/pesos o cualquier otra cosa ocupar más y más memoria?

Creo que probablemente algo esté mal con mi canalización de preprocesamiento de datos, pero he escrito intencionalmente funciones de preprocesamiento como generadores. ¿Podría ser algún tipo de devolución de llamada predeterminada de Keras aplicada al modelo que realiza un seguimiento de la información de entrenamiento responsable del aumento del uso de la memoria?

Mostrar la mejor respuesta

El campo VIRT no importa mucho. Debe comprobar los campos RES o %MEM para conocer la cantidad de memoria utilizada.

@ Yu-Yang Correcto, de acuerdo. Ahora me he dado cuenta de que el consumo de memoria real aumenta durante el proceso de entrenamiento y no durante la inicialización.

Supongo que he encontrado la raíz del problema. Efectivamente, no tenía nada que ver con tensorflow o keras, sino con mi enfoque para usarlos.

Aquí hay una función similar a mi función de aumento de imágenes:

def augment_images():
    transformer = ImageDataGenerator()
    while True:
        x, y = yield
        generator = transformer.flow(x, y, batch_size=len(x), shuffle=False)
        transformed = next(generator)
        yield transformed

Utiliza la clase keras.preprocessing.image.ImageDataGenerator para aumentar las imágenes. Pero esa clase en sí misma instancia el objeto NumpyArrayIterator que mantiene referencias a x y y lotes y llama a ImageDataGenerator como delegar. Y esa fue la fuente de la fuga de memoria. Parece que estos objetos impedían que los arreglos fueran recolectados como basura.

Aquí hay una función de aumento actualizada que usa el iterador explícitamente:

def augment_images(width_shift=0.2,
                   height_shift=0.2,
                   zoom=0.2,
                   rotation=30,
                   vertical_flip=False,
                   horizontal_flip=False):

    transformer = ImageDataGenerator()
    iterator = None

    while True:
        x, y = yield
        if iterator is None:
            iterator = NumpyArrayIterator(
                x, y, transformer,
                batch_size=len(x),
                shuffle=False,
                seed=None,
                data_format=transformer.data_format)
        else:
            iterator.n = x.shape[0]
            iterator.x = x
            iterator.y = y
        transformed = next(iterator)
        yield transformed

Entonces, el problema estaba en los envoltorios del generador que estaba usando para preprocesar los datos. (O diría, en mi método de usar la API de Keras y los generadores de Python). Al menos ahora, cuando reemplacé la función de aumento de imagen, ya no hay pérdidas de memoria.