mesh
mesh copied to clipboard
Ability to add Custom Tensorflow Hooks
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.