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.