Repository: tensorflow/model-analysis
Issue: Model Analysis is very slow.
Root cause: Algorithmic Error

--------------------------------------------------------------------------------

### System information

-   **Have I written custom code (as opposed to using a stock example script
    provided in TensorFlow Model Analysis)**: Yes
-   **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Debian 10 
-   **TensorFlow Model Analysis version (use command below)**: 0.26.0
-   **Python version**: 3.7.9
-   **Jupyter Notebook version**: 6.1.6
-   **Exact command to reproduce**: `tfma.run_model_analysis` or `context.run(evaluator)` (TFX standard component)

Running in an AI Platform Notebook with 4 CPU cores and 15 GB of RAM using the interactive context. 

### Describe the problem

Tensorflow Model analysis takes a very long time to run. For context, I am try to run evaluation on a very small dataset (only 8,000 examples) generated by an example_gen component in a tfx pipeline. If I manually create a tfrecord dataset and run `model.evaluate` (its a Keras model) on it, it will only take a minute or two, including the time it takes to load the model. Using TFMA (either directly or with the evaluator component) takes about 9 minutes, and if I provide a slicing key (with only 6 unique values) it takes 16 minutes. 

I am not using the dataflow runner and I understand that running on dataflow will allow me to scale out my evaluation, but that seems like complete overkill for just a few thousand rows of data. Is it expected that running TFMA should create this much overhead? Is there anything I might be doing that's making this slower than it could be/is there anything I could try to make it go faster? I'm currently only calculating a single metric (in addition to the two saved with my Keras model). 

I'm really liking the capabilities of TFMA, but its currently slowing things down quite a bit.

### Source code / logs

As you can see, my configuration is very simple: 

```python
accuracy_threshold = tfma.MetricThreshold(
    value_threshold=tfma.GenericValueThreshold(
        lower_bound={'value': 0.0001}
    )
)

metrics_spec = tfma.MetricsSpec(
    metrics = [
        tfma.MetricConfig(class_name='BinaryCrossentropy', threshold=accuracy_threshold)
    ]
)

eval_config = tfma.EvalConfig(
    model_specs=[
        tfma.ModelSpec(label_key='series_ep_tags_xf', signature_name="serving_default", preprocessing_function_names=['tft_layer'])
    ],
    metrics_specs=[metrics_spec],
    slicing_specs=[tfma.SlicingSpec()]
)

evaluator = Evaluator(
    examples=example_gen.outputs['examples'],
    model=trainer.outputs['model'], 
    eval_config=eval_config
)
```


Just wanted to give an update that I've now tried this Kubeflow Pipelines as well. Adding in the evaluator component with the eval config above makes pipeline execution go from ~25 minutes to ~40 minutes, which makes it basically unusable given i'm only evaluating on ~10,000 examples currently. 

Prior to using the evaluator component, I was using a stand-in evaluator component that simply loaded in the model and evaluated on the evaluation dataset, roughly like so: 

```python
loaded_model = load_model(os.path.join(trained_model.uri, "serving_model_dir"))

tf_transform_output = tft.TFTransformOutput(transform_graph.uri)
num_tags = tf_transform_output.vocabulary_size_by_name("tags")
tag_file = tf_transform_output.vocabulary_file_by_name("tags")

eval_files = os.path.join(transform_examples.uri, "eval/*")
table = component_utils.create_tag_lookup_table(tag_file)

eval_dataset = component_utils._input_fn(
    file_pattern=eval_files,
    tf_transform_output=tf_transform_output,
    batch_size=64,
    epochs=1,
    num_tags=num_tags,
    table=table,
    shuffle=False,
)
_, precision, recall = loaded_model.evaluate(eval_dataset)

if (precision >= precision_threshold) and (recall >= recall_threshold):
    io_utils.write_string_file(os.path.join(blessing.uri, "BLESSED"), "")
    blessing.set_int_custom_property("blessed", 1)
else:
    io_utils.write_string_file(os.path.join(blessing.uri, "NOT_BLESSED"), "")
    blessing.set_int_custom_property("blessed", 0)

return {"precision": precision, "recall": recall}
```

I was doing this because I was transforming labels in the transform component and as far as I know it wasn't possible to use the standard TFX evaluator component with transformed labels until TFX/TFMA version 0.26.0. This usually took about a minute or so. I really like the TFMA capabilities and the visualisations, but its sort of hard for me to justify its use if pipeline execution time will nearly double, and that's with a very simple configuration. Any help on this would be appreciated! 

@tgaddy 

Could you please upgrade to latest version of TFMA 0.37.0 and let us know if the issue still persists in newer version. Thanks!

@tgaddy 

Closing this issue duo to inactivity. Please feel free to reopen if this still persists. Thanks