Custom eval hooks

tensor2tensor/data_generators/problem.py

@@ -229,6 +229,8 @@ class Problem(object):
   Eval:
     * eval_metrics
         - Specify the set of evaluation metrics for this problem.
+    * eval_hooks
+        - Specify the set of evalueation hooks for this problem.
 
   Inference:
     * feature_encoders(data_dir)
@@ -365,6 +367,9 @@ def eval_metrics(self):
         metrics.Metrics.ACC_PER_SEQ, metrics.Metrics.NEG_LOG_PERPLEXITY
     ]
 
+  def eval_hooks(self, features, logits, hparams):
+    return []
+
   @property
   def task_id(self):
     if self._task_id == -1 and hasattr(self, "global_task_id"):

tensor2tensor/utils/t2t_model.py

@@ -1451,6 +1451,7 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
       raise NotImplementedError(_no_problem_err("estimator_spec_eval"))
 
     problem = hparams.problem
+
     if common_layers.is_xla_compiled():
       remove_summaries()
       if isinstance(logits, dict):
@@ -1509,6 +1510,8 @@ def estimator_spec_eval(self, features, logits, labels, loss, losses_dict):
           summary_op=tf.summary.merge_all())
       evaluation_hooks.append(eval_summary_hook)
 
+      evaluation_hooks += problem.eval_hooks(features, logits, hparams)
+
       return tf.estimator.EstimatorSpec(
           tf.estimator.ModeKeys.EVAL,
           predictions=predictions,