tensorflow/mesh

Ability to add Custom Tensorflow Hooks

trisongz opened this issue · 0 comments

Will there be any future plans to allow users to add Custom Tensorflow Hooks such as tf.estimator.LoggingTensorHook to enable custom functions during the training/eval loop such as passing back metrics to 3rd Party Services

When constructing the TPU estimator model mesh_tensorflow/transformer/utils:tpu_estimator_model_fn, it's fairly complicated to override the training_hooks and evaluation_hooks that is passed to tpu_estimator.TPUEstimatorSpec

Currently, the method I'm using is to override mesh_tensorflow.ops.MtfCheckpointSaverListener functions to enable custom logging/monitoring during training, which doesn't get called as frequently to capture training loss.