Tfx: Passing custom parameters to Trainer and Transform

[rclough]

A common need from TFX users in my organization (using TFX on KFP) is to have the ability to pass custom parameters to their module files for TFT and Trainer to make iteration and experimentation easier. While the Trainer component has hparams in the trainer_fn you use, there doesn't seem to be support for adding anything outside of the TFX DSL (for example, if you wanted to change num_dnn_layers as in the public taxi example's module_file).

Similarly for Transform, there's a number of parameters that users would like to adjust as they explore/compare experiment runs, but there is no means to do this, other than to change them in the module_file, and re-compile the pipeline with the new file each time (as opposed to say, adjusting these parameters through the KFP UI so that you can compare runs, or otherwise do hyper parameter tuning easier).

To support this, we've basically created copy-paste versions of the TFX components in our own internal repo, with a little bit of additional code to support passing these parameters as component-level inputs (dict inputs that are then added to the hparams object, in the case of transform, adding hparams to the preprocessing_fn). This leads to some difficulty with keeping them up to date as TFX releases new versions/fixes.

Which leads me to my question, is this something TFX would incorporate if we were to submit a PR? If not, what is the philosophy around not including such functionality?

Appendix, example code of how it functions:

In your component specification:

from internal_lib import custom_components

...

trainer = custom_components.Trainer(
      module_file=taxi_pipeline_utils,
      train_files=transform_training.outputs.output,
      eval_files=transform_eval.outputs.output,
      schema=infer_schema.outputs.output,
      tf_transform_dir=transform_training.outputs.output,
      train_steps=10000,
      eval_steps=5000,
      warm_starting=True,
      hparams=dict(hidden_layers=6)
      )

Later, in the module_file

def trainer_fn(hparams, schema):
    # hparams can now reference user defined variables from the input dict
    first_dnn_layer_size = 100
    num_dnn_layers = hparams.hidden_layers
    dnn_decay_factor = 0.7

...

    # later, hparams also has all the previously expected values
    estimator = _build_estimator(
      tf_transform_dir=hparams.tf_transform_dir,

      # Construct layers sizes with exponetial decay
      hidden_units=[
          max(2, int(first_dnn_layer_size * dnn_decay_factor**i))
          for i in range(num_dnn_layers)
      ],
      config=run_config,
      warm_start_from=hparams.warm_start_from)

...

Similarly, for Transform (only including 3 hparams for demonstration/readability purposes):

transform_training = components.Transform(
    input_data=examples_gen.outputs.training_examples,
    schema=infer_schema.outputs.output,
    module_file=taxi_pipeline_utils,
    name='transform-training',
    hparams=dict(
        VOCAB_FEATURE_KEYS=["payment_type", "company"], 
        VOCAB_SIZE=1000,
        OOV_SIZE=10
    )
)

And then used in the transform module_file:

def preprocessing_fn(inputs, hparams=None):

...

    for key in hparams.VOCAB_FEATURE_KEYS:
        # Build a vocabulary for this feature.
        outputs[transformed_name(key)] = tft.compute_and_apply_vocabulary(
            fill_in_missing(inputs[key]),
            top_k=hparams.VOCAB_SIZE,
            num_oov_buckets=hparams.OOV_SIZE,
        )

...

[numerology]

I think this is very relevant to #362 . While we are working towards pipeline paramterization, it is definitely a blocker that some TFX components are currently accepting pb as their arguments.
cc @ruoyu90

[rclough]

What we have for the trainer is a relatively simple implementation in the executor, for the user-supplied hparams, we simply grab them from the exec_properties:
python user_hparams = exec_properties.get("hparams", "{}")

and pass them into the hparam constructor here as **user_hparams. For transform there's just changes needed to add hparams as an exec property just as in Trainer, but in this case they would be entirely comprised of the user's input dict.

[ruoyu90]

Thanks @rclough and @numerology ! Actually there is a custom_config field in the trainer interface that can be used for passing more parameters. We can also make the change to pass through the config to the user function. For transform we do not have such flexibility but we can definitely add one. Let me know if this works for you :)

[rclough]

The docs claim this is specifically for CMLE (which we do want to use, but will not always be what we use), but if they could be accessible in the user's trainer_fn that would probably suffice, looks like it would function the same way we use our custom hparams. Also assuming CMLE wont have issues unexpected config values. Same solution would then suffice for transform.

[ruoyu90]

@rclough we'll update the doc to reflect that is not used specifically by CMLE. Will keep this thread opened as a FR for that.

[rclough]

I noticed the copybara PR #699 is blocked on CLA signing (strange thing to ask a bot to do 😛), anything left on that?

Additionally, just wanted to poke on what you thought of including the same functionality for Transform as mentioned before.

[rclough]

I noticed that Transform got a custom_config in the spec in this commit (though I'm not really sure how it was related to the rest of the commit): https://github.com/tensorflow/tfx/commit/1ba6956f2d0fbea2e6f3d1c0b0734c2a59ca371c#diff-51faf1954722449c4cabb94901f737d6L27

But doesn't seem to be used in the component/executor yet

[rclough]

Pinging for an update

[rclough]

Curious how you're thinking it should be exposed for transform- will the preprocessing_fn provide an hparams like trainer does, or just directly pass the config?

[numerology]

After this commit you can specify things in the PARAMETERS section in ComponentSpec as RuntimeParameter, and provide runtime values for them in KFP UI (and it's currently only supported by KubeflowDagRunner). I am working on an example demonstrating its usage.

That being said, there is still some remaining works for a user to let functions in the module file consume those parameters in exec_properties

[rclough]

@numerology thank you for the quick response- that would be handy, though our "internal fork" of TFX already allows that; I'm more concerned with the component level implementations, being able to pass custom arguments to the Trainer (as with #699 ) and similarly to Transform. It sounds like RuntimeParameter will help prioritize this, but the specs for either class will need to have a RuntimeParameter that is passed to the module file to make use of them.

I'm asking mainly because currently, we have copy-pasted versions of Trainer and Transform that only adds a bit of code to pass run time parameters to the module files. We want to avoid having this since it's hard to keep changes up to date with upstream, hence opening this issue.

I'm just trying to align with the planned implementation so that our customers won't see any change when we change from our copy-paste implementation with the upstream implementation once available. I understand if this hasn't been officially decided just yet. We've already moved our internal Trainer to match #699's interface.

[numerology]

@ruoyu90 and @ucdmkt for input re: passing custom_config.

[numerology]

IIUC, executor spec is the encouraged way for the users to bringing their own business logic, so I think it'll be pretty okay if proceed with your own version of trainer and transform executor and expose certain argument of the module files to the component.

[rclough]

I understand that custom business logic should have its own custom executor, but I would argue that this is something that you'd want to adopt as a core part of the component. Being able to pass parameters to your module files at the pipeline level is arguably a common task. Certainly our customers were confused that it wasn't supported out of the box in TFX (even moreso for trainer than tranform)

For example, for transform, you may want to adjust bucket sizes, vocab size, etc through the KFP UI so you can compare their effect on your models accuracy, without having to recompile the pipeline every time (in which case you'd be hard coding them in the module file).

[ruoyu90]

@rclough sorry for the delay, does this commit solve your problem?

[rclough]

Solved for trainer, but also looking for the same functionality for transform

[sngahane]

Hi @ruoyu90, do you have an update for the Transform component?

[rclough]

seems https://github.com/tensorflow/tfx/commit/67b96c016f094398fce9f4b8c760233b55ca47df is addressing this, looks good!

[rclough]

One point I wanted to bring up WRT that commit (adding custom_config to transsform), and something that we've dealt with internally making modified versions of Transform (and other components we have with module_files, which by convention we always include a custom_config so users can pass custom arguments from their pipeline)- have you thought about the consistency of these conventions/interfaces?

We have not found a satisfying answer for consistency, but for example now in Trainer, if you provide a custom_config, all of the first-level keys are now in your train_fn_args or fn_args object, whereas when you use custom_config with Transform, you are explicitly passed a custom_config input in the form of a dict. It's a minor gripe, but can be a bit confusing to end-users.

[ialdencoots]

It would also be nice to be able to pass custom_config in a consistent manner to custom FileBasedExampleGen components. Trainer can take a dict as custom_config while FileBasedExampleGen requires an example_gen_pb2.CustomConfig object which cannot be easily instantiated.