When developing computer vision models, you will often work with datasets that are larger than your device's memory. How do you work with such a dataset, if it can't fit in your memory?

The solution is simply to load the data in chunks (from your hard drive or even the cloud). You load a chunk, process it, and then load the next chunk, leaving the previous ones for the garbage collector to clear. This procedure can be achieved using a generator.

A generator is a programming construct that computes and returns a piece of a whole, one piece at a time, only when the piece is requested. It does not build or load the whole--this might fill up your memory--it, instead, only builds or loads a piece and it does this only when the piece is requested. This makes it memory-efficient and allows you to work with datasets multiple times larger than your memory.

Keras offers an implementation of a generator called ImageDataGenerator. Unfortunately, this generator is designed for datasets where a raw sample is an image file. But I was working on a project where a raw sample was a directory of multiple image files, so ImageDataGenerator could not solve my problem.

Juxtaposition of the dataset structure Keras' ImageDataGenerator works with and the dataset structure my data had — On the left is the dataset structure ImageDataGenerator requires. On the right is my dataset's structure.

Fortunately, Keras offers a utility for building custom generators called Sequence. Sequence is an abstract class that requires the implementation of two methods. The first method should return the number of chunks your dataset is comprised of, whatever size we set our chunks to be. (For an AI dataset, a chunk is a group of n samples and is called a batch while n is called batch-size.) This is the typical implementation of the first method:

(Note that these snippets contain pseudo-code, not actual code, and some things that are not relevant to the subject of the article have been omitted.)

dataset = ...
batchSize = ...

len() -> Number {
  return math.ceil(sizeOf(dataset) / batchSize)
}

The second required method is expected to return a processed batch, given the position of the raw batch (its index) in the dataset. It has two parts; 1. retrieving the raw batch from the dataset and 2. processing the raw batch. It has the following structure:

dataset = ...
batchSize = ...

getItem(idx: Number) -> Whatever {
  # Part 1
  rawBatch = slice(dataset, idx*batchSize, (idx+1) * batchSize)
  
  # Part 2
  return map(rawBatch, (rawSample) => processSample(rawSample, ...))
}

Part 1 is simply a slice operation on the dataset to get a raw batch of batchSize samples starting from the index*batchSize-th sample. Part 2 is a map operation over the retrieved batch using the processSample function to process each sample in the batch. This processSample function was the main challenge of this project.

Processing a Sample

A sample, in my dataset, was a set of images, extracted from a video snippet, maintaining order. In its raw form, a sample was a directory containing the set of images as JPEG or PNG files. Processing a sample meant transforming it from a directory with image files to a 4-dimensional array. The four dimensions were the set of images and the width, height and colour channels of each image in the set.

The first step in processing a sample was getting its image files. For this, the processSample function took a sample's directory name as its first input and would use it to get the sample's image files as an array of their paths.

typealias Sample = Number[][][][]

getSampleImages(sampleDir: Path) -> Path[]
  
processSample(sampleDir: Path) -> Sample {
  imageFiles = getSampleImages(sampleDir)
  //...
}

The second step, after getting the image files, was to standardize the number of image files to be used for all samples. The raw samples had varying numbers of images yet the processed samples had to be formed from a standard user-defined number of images. Therefore, samples that had less than the standard number of images were skipped. On the other hand, for those that had more, this function would systematically pick the required number of images from among its set of images.

Picking the images was not simply taking the first "x" images. That would have been a poor solution because the video snippets from which the samples were derived were of short complete actions therefore just picking the first x images would have resulted in information loss. What I did instead was pick images at intervals through the whole image series - from the start to the end. The interval used for a sample was computed as the division of the number of the sample's images by the required number of images.

Picking images from a sample's image set

pickAtIntervals<T>(list: T[], max: Number) -> T[]
  
processSample(sampleDir: Path, numImages: Number) -> Sample {
  imageFiles = getSampleImages(sampleDir)
  pickedImages = pickAtInterval(imageFiles, numImages)
  //...
}

After picking the required number of images, I would then process the picked images (using Pillow internally) to get them as a 4D array.

processImages(images: Path[]) -> Sample

processSample(sampleDir: Path, numImages: Number) -> Sample {
  imageFiles = getSampleImages(sampleDir)
  pickedImages = pickAtInterval(imageFiles, numImages)
  processedImages = processImages(pickedImages)
  //...
}

The final step was for (optional) image data augmentation. This was done with the help of the excellent Albumentations library which provides image transformation resources. To allow for a user-defined transformation pipeline, the processSample function took a list of Albumentation transformation functions as input; this list would be used to form the transformation pipeline. Importantly, though transformations could be applied randomly across samples, all of a sample's images always had the same transformations applied to them.

transformImages(images: Sample, transforms?: A[]) -> Sample

processSample(
  sampleDir: Path,
  numImages: Number,
  transforms?: A[]
) -> Sample
{
  imageFiles = getSampleImages(sampleDir)
  pickedImages = pickAtInterval(imageFiles, numImages)
  processedImages = processImages(pickedImages)
  return transformImages(processedImages, transforms)
}

After the augmentation step, the sample processing was complete. At this point, I had complete implementations of the two methods required by the Sequence utility and I could leave it to do the rest of the work of loading data in chunks from my hard disk.

Here is a snippet showing the actual generator in use for training an LRCN-like action recognition deep learning model on the Jester dataset

import albumentations as A
from pandasgenerator import PandasGenerator


transformations = [
  A.HorizontalFlip(),
  A.ShiftScaleRotate(shift_limit=0.0625, scale_limit=0.50, rotate_limit=45, p=.75),
  A.Blur(blur_limit=3),
]

train_gen = PandasGenerator(
  source=train_data,
  data_path=data_path,
  nb_frames=NB_FRAMES,
  batch_size=BATCH_SIZE,
  transformations=transformations,
  frame_size=SPATIAL_SIZE
)


model = # An LRCN-like model.
        # input -> MobileNet -> LRCN -> output

model.fit(
  train_gen,
  validation_data=val_gen,
  epochs=30,
  callbacks=[model_checkpoint_callback]
)

Testing

Testing was done using python's unittest framework. Also, a fake programmatically created dataset was used for integration tests on the generator.