小さな画像が大量にある場合、ImageDataGeneratorを使用するとディスクのランダムアクセスネックでGPU使用率が伸びず、学習速度が低下します。学習画像を事前にHDF5に書き込んでおくことで、ランダムアクセスを抑制し、高速化することができます。

まず、flow_from_directoryから取得できる学習用データからh5pyを使用してHDF5を作成します。以下の例ではtrain.h5を作成します。同様にvalidation.h5も作成します。

train_generator = train_datagen.flow_from_directory(
   INPUT_PATH+'/annotations/'+ANNOTATIONS+'/train',
   target_size=(IMAGE_SIZE, IMAGE_SIZE),
   batch_size=BATCH_SIZE,
   class_mode='categorical',
   shuffle=True
)

training_data_n = len(train_generator.filenames)
training_class_n=len(train_generator.class_indices)

import h5py
f = h5py.File("train.h5", 'w')
train_x = f.create_dataset('training_x', (training_data_n,IMAGE_SIZE,IMAGE_SIZE,3), dtype='f')
train_y = f.create_dataset('training_y', (training_data_n,training_class_n), dtype='f')

cnt=0
for x_batch, y_batch in train_generator:
  for i in range(BATCH_SIZE):
    train_x[cnt] = x_batch[i]
    train_y[cnt] = y_batch[i]
    cnt = cnt+1
    if cnt>=training_data_n:
        break
  if cnt>=training_data_n:
    break

f.close()

学習時はHDF5Matrixを使用します。

  from keras.utils.io_utils import HDF5Matrix
  x_train = HDF5Matrix("train.h5", 'training_x')
  y_train = HDF5Matrix("train.h5", 'training_y')
  x_validation = HDF5Matrix("validation.h5", 'validation_x')
  y_validation = HDF5Matrix("validation.h5", 'validation_y')
  fit = model.fit(
    epochs=EPOCS,
    x=x_train, 
    y=y_train,
    validation_data=(x_validation,y_validation),
    batch_size=BATCH_SIZE,
    shuffle='batch'
  )

VGGFace2を使用した性別推定モデル(3146003画像)では160GBのHDF5を生成し、学習時間を1EPOCで3時間から30分に削減でき、6倍高速化できました。評価環境はGeforceRTX2080、32GBメモリ、1TB SSDです。

floatではなく8bit unsigned intでhdf5ファイルに書き込むこともできるようですが、以下のコードで試してみたところ、学習がほとんど進まなくなってしまいました。

train_x = f.create_dataset('training_x', (training_data_n,IMAGE_SIZE,IMAGE_SIZE,3), dtype='uint8')
train_x[cnt] = x_batch[i].astype(np.uint8)

学習時のノーマライズは下記のように行いましたが、どこかに問題があるようです。

ds_norm = lambda x: x / 255.0
x_train = HDF5Matrix(HDF5_PATH, 'training_x', normalizer=ds_norm)