Tuning the batch size in the Tuner component

Hello,

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 hyperparameters (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?

Thanks for your help !

[1025KB]

Customize kerastuner.BaseTuner should work for batch size tuning

You’d pass an unbatched dataset to the Tuner, and in the run_trial method of the CustomTuner, you’d batch the dataset with a variable size drawn from hp:

1. CustomTuner with a run_trial
2. In run trial, get the batch_size from trial
3. Update dataset with correct batch_size
4. train the model with updated dataset
5. you might need to update how the metrics is calculated if you want to add execution time as one of the evaluation method for trial
6. in your tuner_fn you need to change the fit_args to match the run_trial params (pass in unbatched dataset)

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No