Bad tensor shape when using tensorflow drop_remainder - tensorflow

Question

I am trying to create a tf.data.Dataset from a generator. I want to make sure all of my batches have the exact same size, so I'm calling .batch(batch_size, drop_remainder=True) on my Dataset. Here's the relevant code:
train_data = tf.data.Dataset.from_generator(
lambda: map(tuple, train_generator),
(tf.float32, tf.float32),
(
tf.TensorShape([batch_size, crop_height, crop_width, 3]),
tf.TensorShape([batch_size, crop_height, crop_width, 3]),
),
)
val_data = tf.data.Dataset.from_generator(
lambda: map(tuple, val_generator),
(tf.float32, tf.float32),
(
tf.TensorShape([batch_size, crop_height, crop_width, 3]),
tf.TensorShape([batch_size, crop_height, crop_width, 3]),
),
)
my_train_data = train_data.batch(batch_size, drop_remainder=True)
my_val_data = val_data.batch(batch_size, drop_remainder=True)
But I get this error when I run it:
tensorflow.python.framework.errors_impl.InvalidArgumentError: input must be 4-dimensional[4,4,64,64,48] [Op:FusedBatchNormV3]
I get this error because I'm batching the data twice (batch_size is 4 in my error message). I tried to replace the batch_size with None in the .from_generator command, but I get the same error. If I remove the first argument completely like so:
(tf.TensorShape([options["crop_height"], options["crop_width"], 3]),
tf.TensorShape([options["crop_height"], options["crop_width"], 3]),
)
I get this error:
ValueError: `generator` yielded an element of shape (4, 128, 128, 3) where an element of shape (128, 128, 3) was expected.
How can I use drop_remainder without batching the data twice?
EDIT:
Adding code associated with generators:
class BaseGenerator(Sequence):
def __init__(
self,
image_filenames,
label_filenames,
batch_size=1,
is_train=True,
preprocess=None,
augment=None,
height=128,
width=128,
shuffle=False,
):
self.indices = np.arange(0, len(image_filenames))
self.image_filenames = np.array(image_filenames)
self.label_filenames = np.array(label_filenames)
self.batch_size = batch_size
self.is_train = is_train
self.preprocess = preprocess
self.augment = augment
self.crop_height = height
self.crop_width = width
self.shuffle = shuffle
self.on_epoch_end() # shuffle data
def __len__(self):
return int(np.ceil(len(self.indices) / float(self.batch_size)))
def __getitem__(self, index):
min_index = index * self.batch_size
max_index = min((index + 1) * self.batch_size, len(self.indices))
batch_indices = self.indices[min_index:max_index]
return self.generate(self.image_filenames[batch_indices], self.label_filenames[batch_indices])
def __call__(self):
return next(iter(self))
def on_epoch_end(self):
if self.is_train and self.shuffle:
np.random.shuffle(self.indices)
def generate(self, image_filenames, label_filenames):
X = np.zeros((self.batch_size, self.crop_height, self.crop_width, 3), dtype=np.float32)
y = np.zeros((self.batch_size, self.crop_height, self.crop_width), dtype=np.float32,)
for i, (image_fn, label_fn) in enumerate(zip(image_filenames, label_filenames)):
image = utils.load_image(image_fn)
label = utils.load_image(label_fn)
if self.augment:
augmented = self.augment(image=image, mask=label)
image = augmented["image"]
label = augmented["mask"]
if self.preprocess:
image = self.preprocess(image)
label = np.float32(helpers.one_hot_it(label=label))
X[i, :, :, :] = image
y[i, :, :, :] = label
return X, y
train_generator = BaseGenerator(
image_filenames=train_input_names,
label_filenames=train_output_names,
batch_size=batch_size,
is_train=True,
preprocess=preprocessing,
augment=None,
height=128,
width=128,
)
val_generator = BaseGenerator(
image_filenames=val_input_names,
label_filenames=val_output_names,
batch_size=batch_size,
is_train=False,
preprocess=preprocessing,
augment=None,
height=128,
width=128,
)

Answer 1

As you mentioned in the question, the issue is that you are batching your data twice. To overcome this problem, you can:
First, define a generator that yields single images (e.g. without batch dimension).
Then, group your examples into batches using the method batch of tf.data.Dataset.
In order to redefine BaseGenerator so that it yields single images, you can follow the next steps.
First, in the __init__ method, remove batch_size because it is no longer needed:
def __init__(
self,
image_filenames,
label_filenames,
is_train=True,
preprocess=None,
augment=None,
height=128,
width=128,
shuffle=False,
):
self.indices = np.arange(0, len(image_filenames))
self.image_filenames = np.array(image_filenames)
self.label_filenames = np.array(label_filenames)
self.is_train = is_train
self.preprocess = preprocess
self.augment = augment
self.crop_height = height
self.crop_width = width
self.shuffle = shuffle
self.on_epoch_end() # shuffle data
Second, adapt the method generate so that it yields a single example:
def generate(self, image_filename, label_filename):
image = utils.load_image(image_filename)
label = utils.load_label(label_filename)
if self.augment:
augmented = self.augment(image=image, mask=label)
image = augmented["image"]
label = augmented["mask"]
if self.preprocess:
image = self.preprocess(image)
label = np.float32(helpers.one_hot_it(label=label))
X = image # Shape=(self.crop_height, self.crop_width, 3)
Y = label # Shape=(self.crop_height, self.crop_width)
return X, y
Third, in the method __getitem__, pass only one filename:
def __getitem__(self, index):
return self.generate(self.image_filenames[index], self.label_filenames[index])
Finally, exclude the batch dimension when defining your tf.data.Dataset:
train_data = tf.data.Dataset.from_generator(
lambda: map(tuple, train_generator),
(tf.float32, tf.float32),
(
tf.TensorShape([crop_height, crop_width, 3]),
tf.TensorShape([crop_height, crop_width]),
),
)
my_train_data = train_data.batch(batch_size, drop_remainder=True)
it = iter(my_train_data)
x, y = next(it)
print(x.shape) # (4, 128, 128, 3)
print(y.shape) # (4, 128, 128)