What is an alternative to tf.data.Dataset.from_tensor_slices for a multi-input model?

Viewed 420

I am trying to make a multi-input Keras model that takes two inputs. One input is an image, and the second input is 1D text. I am storing the paths to the images in a dataframe, and then appending the images to a list like this:

from tqdm import tqdm

train_images = []
for image_path in tqdm(train_df['paths']):
  byte_file = tf.io.read_file(image_path)
  img = tf.image.decode_png(byte_file)
  train_images.append(img) 

The 1D text inputs are stored in lists. This process is repeated for the validation and test sets. I then make a dataset, like this:

train_protein = tf.expand_dims(padded_train_protein_encode,axis=2)
training_dataset = tf.data.Dataset.from_tensor_slices(((train_protein, train_images), train_Labels)) 

training_dataset = training_dataset.batch(20)

val_protein = tf.expand_dims(padded_val_protein_encode, axis=2)
validation_dataset = tf.data.Dataset.from_tensor_slices(((val_protein, val_images), validation_Labels))
validation_dataset = validation_dataset.batch(20)

test_protein = tf.expand_dims(padded_test_protein_encode, axis=2)
test_dataset = tf.data.Dataset.from_tensor_slices(((test_protein, test_images), test_Labels)) 
test_dataset = test_dataset.batch(20)

I am running this in Google Colab, and even using the high-ram option, the program crashes due to running out of ram. What is the best way to solve this problem?

I have researched tf.data.Dataset.from_generator as an option, but I can't work out how to make it work when there are two inputs. Can anyone help?

1 Answers

This is a fairly common pain. There really isn't a better way than a datagenerator if your dataset is too large to load into memory. Coming from PyTorch, there are pythonic classes to do this, rather than having to use tf.data.Dataset.from_generator. Subclassing tf.keras.utils.Sequence could be an elegant alternative. Without access to your dataset, I cannot verify but something like this should work.

__getitem__ is called every batch.

class TfDataGenerator(tf.keras.utils.Sequence):
    def __init__(self, filepaths, proteins, labels):
        self.filepaths = np.array(filepaths)
        self.proteins = np.array(proteins)
        self.labels = labels

    def __len__(self):
        return len(self.filenames) // self.batch_size

    def __getitem__(self, index):
        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]
        return __generate_x(indexes), labels[indexes]

    def __generate_x(self, indexes):
        x_1 = np.empty((self.batch_size, *self.dim, self.n_channels))
        x_2 = np.empty((self.batch_size, len(self.meta_features)))

        for index in enumerate(indexes):
            image = cv2.imread(self.filepaths[index])
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            x_1[num] = image.astype(np.float32)/255.
            x_2[num] = self.proteins[index]

        return [x_1, x_2]

    def on_epoch_end(self):
        self.indexes = np.arange(len(self.filenames))
        if self.shuffle:
            np.random.shuffle(self.indexes)

Again, a very rough example, but hopefully it shows what can be done. Tensorflow documentation here

This has been a big headache for me in the past, so hopefully this answer helps.

Related