TFMA не удалось найти метрики для модели Keras при загрузке результата оценки

Question

Я обучил модель Keras (не оценщик) со следующей сигнатурой обслуживания:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['examples'] tensor_info:
        dtype: DT_STRING
        shape: (-1)
        name: serving_default_examples:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['mu'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall_1:0
    outputs['sigma'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall_1:1
  Method name is: tensorflow/serving/predict

Веса обновляются с использованием настраиваемого цикла обучения с градиентной лентой вместо model.fit метод, прежде чем модель будет экспортирована как saved_model. Поскольку я не могу заставить TFMA работать без предварительной компиляции модели, я компилирую модель, указывая набор пользовательских метрик Keras:

model.compile(metrics=custom_keras_metrics) # each custom metric inherits from keras.Metric
custom_training_loop(model)
model.save("path/to/saved_model", save_format="tf")

Я хотел бы оценить эту модель с помощью TFMA, поэтому сначала я инициализирую общую модель eval следующим образом:

eval_config = tfma.EvalConfig(
    model_specs=[tfma.ModelSpec(label_key="my_label_key")],
    slicing_specs=[tfma.SlicingSpec()] # empty slice refers to the entire dataset
)
eval_shared_model = tfma.default_eval_shared_model("path/to/saved_model", eval_config=eval_config)

Однако, когда я пытаюсь запустить анализ модели:

eval_results = tfma.run_model_analysis(
    eval_shared_model=eval_shared_model,
    data_location="path/to/test/tfrecords*",
    file_format="tfrecords"
)

Я столкнулся со следующей ошибкой:

ValueError          Traceback (most recent call last)
<ipython-input-156-f9a9684a6797> in <module>
      2     eval_shared_model=eval_shared_model,
      3     data_location="tfma/test_raw-*",
----> 4     file_format="tfrecords"
      5 )

~/.pyenv/versions/miniconda3-4.3.30/envs/tensorflow/lib/python3.7/site-packages/tensorflow_model_analysis/api/model_eval_lib.py in run_model_analysis(eval_shared_model, eval_config, data_location, file_format, output_path, extractors, evaluators, writers, pipeline_options, slice_spec, write_config, compute_confidence_intervals, min_slice_size, random_seed_for_testing, schema)
   1204 
   1205   if len(eval_config.model_specs) <= 1:
-> 1206     return load_eval_result(output_path)
   1207   else:
   1208     results = []

~/.pyenv/versions/miniconda3-4.3.30/envs/tensorflow/lib/python3.7/site-packages/tensorflow_model_analysis/api/model_eval_lib.py in load_eval_result(output_path, model_name)
    383       metrics_and_plots_serialization.load_and_deserialize_metrics(
    384           path=os.path.join(output_path, constants.METRICS_KEY),
--> 385           model_name=model_name))
    386   plots_proto_list = (
    387       metrics_and_plots_serialization.load_and_deserialize_plots(

~/.pyenv/versions/miniconda3-4.3.30/envs/tensorflow/lib/python3.7/site-packages/tensorflow_model_analysis/writers/metrics_and_plots_serialization.py in load_and_deserialize_metrics(path, model_name)
    180       raise ValueError('Fail to find metrics for model name: %s . '
    181                        'Available model names are [%s]' %
--> 182                        (model_name, ', '.join(keys)))
    183 
    184     result.append((

ValueError: Fail to find metrics for model name: None . Available model names are []

Почему TFMA вызывает это исключение и с чего начать отладку этой ошибки? Я попытался указать названия моделей вручную (в этом нет необходимости, поскольку я использую только одну модель), но это тоже не помогло. Я попытался отследить исходный код, и похоже, что это происходит, когда TFMA пытается загрузить результат eval, сгенерированный PTransform.

Я использую tensorflow==2.3.0 а также tensorflow-model-analysis==0.22.1