How to use TPU on Kaggle/Colab with ImageDataGenerator?

I have a directory and dataframe that looks like this:

|---classA
|-----1.jpg
|-----2.jpg
|---classB
|-----23.jpg
|-----124.jpg

The dataframe consists of two columns: the image's name and its classes

name  | class
1.jpg |   1
2.jpg |   1
23.jpg|   2

I use ImageDataGenerator and use the method flow_from_directory. I am currently using GPU to train but since it is big it takes a lot of time, so I want to switch to TPU. I searched both Kaggle and Colab's TPU guide, and they use TFRecord to read the data. Can anyone help me how to use TFRecord from ImageDataGenerator?

Solution

So I find out that the ImageDataGenerator cannot be used on TPU. You have to use tf.data.Dataset instead.