Tensorflow MirroredStrategy halves the 2nd dimension, though shape in the object remains right

I’ve recently tried to use MirroredStrategy for training. The relevant code is:

dataset = tf.data.Dataset.from_tensor_slices((samples, feature))
model.fit(dataset,
          batch_size=500, epochs=150,
          #callbacks=[tensorboard_callback]
          )

dataset print is:

Out[1]: <TensorSliceDataset element_spec=(TensorSpec(shape=(16, 17, 10), dtype=tf.float64, name=None), TensorSpec(shape=(1000,), dtype=tf.float32, name=None))>

which is in the correct dimension, but I get the following error:

ValueError: Input 0 of layer "sequential" is incompatible with the layer: expected shape=(None, 16, 17, 10), found shape=(8, 17, 10)

which is odd, as the documentation says that the strategy will halve the first dimension not the second, it should split the dataset for 2, along the first axis.

Does anyone know what is the problem?

sidenote: the name of the variable “feature” is misleading, it should be label instead.

Answer

The Keras API, while awesome in many ways, is also long in the tooth and fit in particular is heavily overloadeded.

In your case, don’t use the batch_size arg in fit… batch your dataset instead via

dataset = dataset.batch(500).

see docs for keras fit:

batch_size  Integer or None. Number of samples per gradient update. If unspecified, batch_size will default to 32. Do not specify the batch_size if your data is in the form of datasets, generators, or keras.utils.Sequence instances (since they generate batches).