使用估算器训练自定义Kera模型时如何打印指标值？

Question

我已经使用Model subclass定义了一个自定义的keras模型。我想将此模型转换为用于分布式训练的估计器。简单的代码如下：

    tf.logging.set_verbosity(tf.logging.INFO)

    model = EstimatorModel()

    model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    estimator_model = tf.keras.estimator.model_to_estimator(
        keras_model=model, model_dir=FLAGS.model_path)

    estimator_model.train(input_fn=input_fn)

代码可以正常运行。现在，问题是估计量只打印loss，而没有打印metric。我尝试使用estimator_model将指标添加到tf.estimator.add_metrics(estimator_model, my_auc)，但仍然无法正常工作。我该怎么做才能解决这个问题？