Trial of a DALIDataset with external input for Keras

keras_pipeline_trial.py

import glob
import numpy as np
import cupy as cp
import imageio
from random import shuffle
from nvidia.dali import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.plugin.tf as dali_tf
import tensorflow as tf

BATCH_SIZE = 16
FILE_PATH_GLOB = "/files/data/pokemon_jpg/*.jpg" # replacing with local path

class ExternalInputGpuIterator(object):
    def __init__(self, batch_size):
        self.files = glob.glob(FILE_PATH_GLOB)
        self.batch_size = batch_size
        shuffle(self.files)

    def __iter__(self):
        self.i = 0
        self.n = len(self.files)
        return self

    def __next__(self):
        batch = []
        for _ in range(self.batch_size):
            im = imageio.imread(self.files[self.i])
            im = cp.asarray(im)
            im = im / 255.0
            batch.append(im.astype(cp.float32))
            self.i = (self.i + 1) % self.n
        return (batch,)

# creating and testing iterator
eii_gpu = ExternalInputGpuIterator(BATCH_SIZE)
print(type(next(iter(eii_gpu))[0][0]))

# creating pipeline
pipe_gpu = Pipeline(batch_size=BATCH_SIZE, num_threads=2, device_id=0)
with pipe_gpu:
    images = fn.external_source(source=eii_gpu, num_outputs=1, device="gpu")
    images = fn.resize(images, size=[256,256])
    pipe_gpu.set_outputs(images)
pipe_gpu.build()

# testing pipeline
pipe_out_gpu = pipe_gpu.run()
batch_gpu = pipe_out_gpu[0].as_cpu()
img = batch_gpu.at(0)
print(img.shape)

# trying to create DALIDataset
shapes = ((BATCH_SIZE, 256, 256, 3),)
dtypes = (tf.float32,)

with tf.device('/gpu:0'):
    dataset = dali_tf.experimental.DALIDatasetWithInputs(
        pipeline=pipe_gpu,
        batch_size=BATCH_SIZE,
        output_shapes=shapes,
        output_dtypes=dtypes,
        device_id=0)