keras fit_generator() with hand-built image data generator: out of RAM

Viewed 278

I am writing a Jupyter notebook in Google Colab, which uses keras to fit a pre-trained model to an image classification dataset. Even though I think my handcrafted generator should not lead to memory issues, the RAM is filling up and causes the Colab session to crash. I tried to boil down my code as much as possible (points 4, 5, and 6 are the relevant ones):

1) a few imports

import os
import tarfile
import urllib.request

import numpy as np
import pandas as pd
import skimage.io
import tensorflow as tf

2) dataset preparation

dataset = {
    'name': 'oxford-102-flowers',
    'url': 'https://s3.amazonaws.com/fast-ai-imageclas/oxford-102-flowers.tgz',
    'num_classes': 102,
}

base_dir = os.getcwd()
data_dir = os.path.join(base_dir, 'data')
dataset_dir = os.path.join(data_dir, dataset['name'])
os.makedirs(data_dir)
tgz_file = os.path.join(data_dir, dataset['name']+'.tgz')
urllib.request.urlretrieve(dataset['url'], tgz_file)
tar = tarfile.open(tgz_file, 'r:gz')
tar.extractall(data_dir)
tar.close()
os.remove(tgz_file)

train_txt = os.path.join(dataset_dir, 'train.txt')
train_df = pd.read_csv(train_txt, sep=' ', header=None, names=('file', 'label'))

3) model setup

loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
base_model = tf.keras.applications.MobileNetV2(include_top=False)
pooling_dense_model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(dataset['num_classes']),
])
pooling_dense_model.compile(
    loss=loss,
    optimizer='adam',
    metrics=['acc'],
)

4) build a generator

def make_train_generator():
  while True:
    for i, record in train_df.iterrows():
      filename = os.path.join(dataset_dir, record['file'])
      image = skimage.io.imread(os.path.join(dataset_dir, filename))
      normalised_image = image/255. - 0.5
      image_batch = normalised_image[np.newaxis, ...]
      label_batch = np.array((record['label'],))
      yield image_batch, label_batch

5) check that the generator by itself does not produce memory leaks

train_generator = make_train_generator()
train_len = len(train_df)
for i in range(2*train_len):
  train_example = next(train_generator)

6) train the model

pooling_dense_model_history = pooling_dense_model.fit_generator(
    generator=train_generator,
    steps_per_epoch=train_len,
    epochs=2,
    verbose=1,
)

Now point 5) does not increase RAM consumption, while during the steps of 6) memory pressure increases, until it gets an extra boost at the end of the epoch and the Colab session crashes.

I browsed extensively for solutions related to fit_generator memory consumption, but it seems like all of them have additional complications compared to the setup above.

Any ideas of where I am going wrong?

PS I am using a python3 session with GPU acceleration and tf.__version__ is 1.15.0-rc3.

0 Answers
Related