Tfx: Automatic Hyperparameter Tuning in TFX

[andrewlarimer]

Per @zhitaoli's suggestion a few days ago here, I am opening this up as a new thread.

I am wondering what is the best way to search for a model's optimal hyperparameters within TFX.

It sounds like it may be possible to use the Trainer's executor_class and specify an AI Platform Executor and even pass the additional arguments it needs for it to do hyperparameter tuning, but from Zhitao's answer it sounds like receiving the responses may not be set up yet.

Anyway, I'm curious to hear what the roadmap to hyperparameter tuning within TFX looks like, and if there's anything I can do to contribute to that effort.

[zhitaoli]

Hi @andrewlarimer, I had some discussion with our internal experts on hyper parameter tuning, and the following is what I think:

1. We had a pre-existing internal component for hyper-parameter tuning (and some other tuning support), but it's going through some changes (especially related to tf.Keras) and we don't feel ready to publish something with long term support immediately;
2. If a stop-gap solution as a custom component is useful, we can provide guides on how to achieve this using the Google Cloud API Platform training job. However, it will not be considered as supported solution from TFX;

Please let me know what you think.

Thanks.

[andrewlarimer]

Learning more about custom components with hyperparameter tuning as a specific example would be excellent, and would definitely make me feel more comfortable pitching TFX as a solution to clients, knowing we could likely build other custom components as needed.

So yes that would be very useful in the meantime until the official component is ready for primetime again. I imagine others would find this useful as well. Thank you.

[Efaq]

@zhitaoli I agree with @andrewlarimer, it would be useful to get some insights from you in the meantime about how to implement something through custom components.

Thank you!

[sseveran]

Completely agree. If there was support for katib or preferably ray that would be great. A potential starting point might just be some docs describing how this should be implemented and we could add some support for what we are using today.

[zhitaoli]

Sorry for the long delay as we have a couple of discussions with various teams (Google Cloud AI Platform, TFX internal, Keras, etc).

My current thought process on this one:

1. Given that Tensorflow is moving towards TF 2.0 and Keras, we hope to utilize this window to find a unified programming interface for hyper-parameter related business. The new keras-tuner project seems to provide a good starting point;
2. Once we agree on API to modeling work, different tuner implementations integrated with our pipeline system will be developed by collaboration between TFX::OSS and Cloud AI Platform/Katib/Keras, to realize hyper-parameter tuning on all thees platforms.
3. One important requirement is that the same model code and tuner component configuration should be portable as much as possible to different platform/environment. For instance, moving a locally loop-based tuner to GCP/Kubeflow should only require minimal new options which are only applicable to the specific environments.

@1025KB and I will publish more doc(s) and discussions in this thread as we make progress.

[sseveran]

We would be happy with that if it also supports estimators. There are some still architectures we have not figured out how to create with Keras.

[jinnovation]

Given that Tensorflow is moving towards TF 2.0 and Keras, we hope to utilize this window to find a unified programming interface for hyper-parameter related business.

@zhitaoli: Am I correct, then, in assuming that the current Estimator-centric Trainer Component would likely be updated in the near future to be Keras-centric instead?

[zhitaoli]

@jinnovation Yes we will provide a Keras-centric trainer too, but the actual timeline is also being worked as part of whole TF 2.0 support in TFX. We hope to communicate the phases of such support later this quarter. Please stay tuned.

[schmidt-jake]

@zhitaoli any update on this? Thanks

[1025KB]

Hi, Sorry for the late response, we are working on a tuner prototype, and we will also provide an e2e example after that.

this is currectly under examples/ folder, after we confirmed a stable interface, we will move it to components folder

[rclough]

@1025KB I was looking through the tuner prototype, and had a quick question/clarification- would this tuner essentially be a single threaded HP search? We're currently working on adding support for AI platform HP tuning support in our own (fairly diverged) internal "fork" of TFX, and its a bit painful but we're making progress, but we're also obviously trying to keep an eye on the official HP tuning implementations coming out.

[1025KB]

There are several level of distribution
1) multiple models (with different hyperparameter) tuning in parallel (multiple KerasTuners with a central hp server)
2) for a single model, data parallelism training (single KerasTuner)
3) for a single model, model parallelism training (useful when model is too big to fit into memory, but we will mainly working on 1 and 2)

This prototype tuner is single KerasTuner, multiple hyperparameter combinations will train in sequential instead of in parallel, but we will consider making it data parallel (Op2), we will provide a Google Cloud tuner later with distributed tuning (Op1).

What kind of distribution are you working on?

[rclough]

@1025KB Thanks for the clarification on what's here and what's planned!

We're essentially just working on modifying the existing Trainer component, so that when you use the AI platform executor like this, you can add hyperparameters into the ai_platform_training_args in the custom_config as specified in the AI platform TrainingInput spec.

By default this will not work in the public implementation since HP tuning on AI platform has a number of differences from regular training- in particular it expects to pass the Hyperparamers as CLI args, and then the trial results need to be parsed at the end to select one model to output so that it can conform to the way the existing Trainer works.

We're doing this by passing up the command line arguments from the component runner and using the HP spec in the ai_platform_training_args to look through these args and add them to the custom_config at run time, as we pass custom_config into the trainer_fn similar to how pr #699 does it.

There's a couple of other details I'm missing but this is the broad stroke of our approach. This way we get any level of parallelization that AI platform provides.

[rclough]

We're considering making this a distinct component though, especially due to the special post-processing of the HP trial results....we try to avoid having copy-pasted code from TFX just to add special functionality since its hard to keep in sync, so keeping it separate from Trainer allows us to continue using the default TFX executor, etc. It's not immediately clear to us whether our customers would prefer having a single component with HP tuning on/off vs needing to explicitly switch components at different cycles of their pipeline development

[zhitaoli]

@andrewlarimer
Some update: @1025KB has provided a custom component 'tuner' in upcoming release of TFX. Please see this example on how to use it.

For people who are interested, please feel free to try it out and provide some feedback.

Thanks.

[ucdmkt]

\cc @ucdmkt @yinghsienwu

[1025KB]

Hi, we added a new Tuner component for model hparams tuning