using tf.data for fit method in DeepEnsemble model

[hstojic]

this simple change should improve memory handling, should be better optimized for GPUs, and generally gives more control to the user over preparing data for training

[uri-granta]

Various comments/comments. Happy to review again (including the tests) once the tests are passing.

[uri-granta]

might be nicer to always return a tuple?

Suggested change

	) -> Union[tf.data.Dataset, tuple[tf.data.Dataset, tf.data.Dataset]]:
	) -> tuple[tf.data.Dataset, Optional[tf.data.Dataset]]]:

[uri-granta]

Suggested change

	raise ValueError("validation_split must be between 0 and 1")
	raise ValueError(f"validation_split must be between 0 and 1: got {validation_split}")

[uri-granta]

Suggested change

	val_size = int(num_points * validation_split)
	val_size = round(num_points * validation_split)

[uri-granta]

"batch_size" isn't guaranteed to exist for a user-supplied fit_args

Suggested change

	batch_size=fit_args_copy["batch_size"],
	batch_size=fit_args_copy.get("batch_size"),

[hstojic]

ah, well spotted, I was remembering BatchOptimizer...

[uri-granta]

(if you change the return type above as suggested)

Suggested change

	tf_data = self.prepare_tf_data(
	train_dataset, val_dataset = self.prepare_tf_data(

[uri-granta]

should we maybe raise an exception if "train_dataset, val_dataset" is already present in the fit_args?

[uri-granta]

Suggested change

	history = self.model.fit(
	train_dataset, **fit_args_copy, initial_epoch=self._absolute_epochs
	)
	else:
	history = self.model.fit(tf_data, **fit_args_copy, initial_epoch=self._absolute_epochs)
	history = self.model.fit(tf_data, **fit_args_copy, initial_epoch=self._absolute_epochs)

[uri-granta]

(I think?)

Suggested change

	.shuffle(train_size, reshuffle_each_iteration=True)
	.shuffle(num_points, reshuffle_each_iteration=True)

[uri-granta]

Q: is this really the same as the original behaviour?

[pio-neil]

I do have a small concern about the use of tf.Dataset.shuffle. When I did some testing with this before, the shuffle buffer (which is tf.Dataset's internal method of shuffling data) used around 18GB of extra memory. This was with a dataset with 30 million rows, with 15 inputs and one output, and a batch size of 1000. The shuffle buffer also has an impact on speed, but I suspect this is relatively minor compared to the model training time.

This might not be such a problem with smaller datasets. So perhaps it would be a good idea to do some benchmarking?

However, it's also not clear to me why we're introducing shuffling by default here, when AFAICT it wasn't there before? This seems like a change of behaviour. Do we expect this to improve model accuracy? It may be better to let the user of Trieste control this, rather than making it the default behaviour?

[hstojic]

I do have a small concern about the use of tf.Dataset.shuffle. When I did some testing with this before, the shuffle buffer (which is tf.Dataset's internal method of shuffling data) used around 18GB of extra memory. This was with a dataset with 30 million rows, with 15 inputs and one output, and a batch size of 1000. The shuffle buffer also has an impact on speed, but I suspect this is relatively minor compared to the model training time.

This might not be such a problem with smaller datasets. So perhaps it would be a good idea to do some benchmarking?

I have already done some testing, on much smaller data than that, haven't seen any adverse effects.
Users that handle bigger datasets now can handle data preparation themselves by subclassing the model and overwriting the new method, if there are any issues with speed/memory.

However, it's also not clear to me why we're introducing shuffling by default here, when AFAICT it wasn't there before? This seems like a change of behaviour. Do we expect this to improve model accuracy? It may be better to let the user of Trieste control this, rather than making it the default behaviour?

shuffle argument is by default True in Keras' fit method, hence why its default here, but its definitely needed - we are typically doing many epochs here and shuffling data at each epoch helps model training - so yes, should improve model accuracy. The user can always overwrite the new method, that was the purpose of extracting this data preparation bit into a method.

[pio-neil]

However, it's also not clear to me why we're introducing shuffling by default here, when AFAICT it wasn't there before? This seems like a change of behaviour. Do we expect this to improve model accuracy? It may be better to let the user of Trieste control this, rather than making it the default behaviour?

shuffle argument is by default True in Keras' fit method, hence why its default here, but its definitely needed - we are typically doing many epochs here and shuffling data at each epoch helps model training - so yes, should improve model accuracy. The user can always overwrite the new method, that was the purpose of extracting this data preparation bit into a method.

Right, yes, I see. I was just wondering why we were doing things differently in terms of shuffling, but we aren't (since fit ignores the shuffle argument for datasets):

shuffle: Boolean, whether to shuffle the training data before each epoch. This argument is ignored when x is a keras.utils.PyDataset, tf.data.Dataset, torch.utils.data.DataLoader or Python generator function.

[hstojic]

after some further thinking, this is currently unnecessary complication of the code, as for smallish datasets that we mainly deal with here current code is good enough