I'm currently using TFX to build a pipeline on the Google AI platform with the Kubeflow engine. I have a model where the batch size is an important hyper-parameter to tune.
I would like to search this hyper-parameter in the Tuner component.
Is it even possible?
I follow the TFX example with the Penguin dataset, more precisely the tuner component implementation: found here.
The _get_hyperparameters function returns the sample space for the model hyper-parameters (see line 139).
However, the batch size to train the model is fixed and specified at the end of the tuner_fn (see line 246).
Is there a way to dynamically change the batch size based on a sample from the hyper-parameter space?