'Keras - How to utilize GPU with custom generator

I have developed a custom generator to generate batches for my deep learning model to use when training. Its purpose is to minimize the amount of memory being used while providing batches of augmented images on the fly. The augmentation methods are all contained in the ImageAugmentor_v2.py file and make use of various OpenCV, Pillow, and NumPy methods. Naturally, when starting the training, the process is very slow, advancing at about 1 step per second, suggesting that the CPU is doing all the work. After having performed a couple tests, I figured out that the ImageAugmentor_v2 class is the culprit, slowing everything down.

Given the fact that I am using all of the aforementioned libraries' high-level methods, is there still a way to perform all of the computation on the GPU? I have looked into the various tools out there such as CuPy and Numba but I am not sure how they can be used in the particular class.

CustomDataGenerator.py

import tensorflow as tf
from tools.ImageAugmentor_v2 import ImageAugmentorV2
import numpy as np

class CustomDataGenerator(tf.keras.utils.Sequence):
    
    def __init__(self, X, Y, batch_size,
                 new_size_multiplier,
                 aug_params,
                 input_size=(128,128,3),
                 shuffle=True):
        
        self.X = X
        self.Y = Y
        self.batch_size = batch_size
        self.input_size = input_size
        self.shuffle = shuffle
        self.indexes = np.arange(len(self.X))
        
        self.augmentor = ImageAugmentorV2(**aug_params)
        
        self.original_size = len(self.X)
        
        self.new_size = len(self.X)*new_size_multiplier
        
    def on_epoch_end(self):
        pass
    
    def __getitem__(self, index):
        batches = self.indexes[(index * self.batch_size)%self.original_size:((index + 1) * self.batch_size)%self.original_size]
        
        x_batch = self.X[batches]
        y_batch = self.Y[batches]
        
        X, y = self.augmentor.flow(x_batch, y_batch, combine_augmentations=True, verbose=0)      
        return X, y[..., np.newaxis].astype(bool)
    
    def __len__(self):
        return self.new_size // self.batch_size


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source