How should metrics be added to a multi-headed TensorFlow estimator?
Build a class around the model creation that holds the model configuration and use a member function for the metric_fn parameter.
class ModelBuilder:
  # constructor storing configuration options in self
  def __init__(self, labels, other_config_args):
    self.labels = labels
    ...
  # Function for building the estimator with multiple heads (multi-objective)
  def build_estimator(self, func_args):
    heads = []
    for label in self.labels:
      heads.append(tf.estimator.MultiClassHead(n_classes=self.nclasses[label], name=label))
    head = tf.estimator.MultiHead(heads)
    estimator = tf.estimator.DNNEstimator(head=heads,...) # or whatever type of estimator you want
    estimator = tf.estimator.add_metrics(estimator, self.model_metrics)
    return estimator
  # Member function that adds metrics to the estimator based on the model configuration params
  def model_metrics(self, labels, predictions, features):
    metrics = {}
    for label in self.labels: # generate a metric for each head name
      metrics['metric_name'] = metric_func(features,labels,predictions[(label,'logits')])
    return metrics
I previously created metrics for a TensorFlow Classifier referencing predictions['logits'] to calculate the metrics. I have changed the model from a Classifier to an Estimator in order to enable multi-objective learning (using MultiHead). However, this has caused Python to throw an error since now the elements of predictions are keyed by pairs of the head name and original key, e.g. ('label1','logits') for a head with name 'label1'.
I'd like to allow for dynamic generation of metrics based on a configuration file in order to more easily train and test a variety of models with different label combinations. The problem now is that the metric_fn parameter to tf.estimator.add_metrics does not take any additional parameters to allow for dynamically determined or constructed metrics.
How can I generate an estimator with multiple heads and custom metrics for each head?