keras fit_generator () со встроенным генератором данных изображения: недостаточно ОЗУ

Question

Я пишу блокнот Jupyter в Google Colab, в котором используется keras для подгонки предварительно обученной модели к набору данных классификации изображений. Хотя я думаю, что мой генератор ручной работы не должен приводить к проблемам с памятью, ОЗУ заполняется и вызывает сбой сеанса Colab. Я пытался свести мой код как можно больше (пункты 4, 5 и 6 являются соответствующими):

1) несколько импортов

import os
import tarfile
import urllib.request

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

2) подготовка набора данных

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) настройка модели

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) сборка генератора

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) проверка того, что сам генератор не вызывает утечек памяти

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

6) обучаем модель

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

Теперь пункт 5) не увеличивает потребление ОЗУ, в то время как на этапах 6) давление памяти возрастает, пока не получит дополнительный импульс наконец эпохи и сбой сеанса Colab.

Я много раз искал решения, связанные с fit_generator использованием памяти, но, похоже, у всех них есть дополнительные сложности по сравнению с настройкой выше.

Есть идеи, где я ошибаюсь?

PS Я использую сеанс python3 с ускорением графического процессора, а tf.__version__ равен 1.15.0-rc3.