TensorFlow Model Analysis: evaluating a model with the metrics its trainer already defined

Evaluating with the metrics the trainer already defined

TensorFlow Model Analysis is a library for evaluating TensorFlow models, and the README's description of what that means is precise. It allows users to evaluate their models on large amounts of data in a distributed manner, using the same metrics defined in their trainer.

That last clause is the whole idea. Most evaluation codebases end up with two implementations of every metric, one in the trainer and one in the evaluation job, and the two drift apart until the evaluation stops meaning anything. TFMA takes the metric from the model, which is what makes the numbers comparable.

The scale claim, large amounts of data in a distributed manner, is the second half. Evaluation on a big slice of production data is a batch job over a data warehouse, not a notebook, and TFMA is built for that shape of work, with Apache Beam as the execution model underneath the distributed path.

The intended user is a machine learning engineer at a company with a trained model in production and a question that the training loss does not answer, usually of the form why is this worse for this segment. The third piece of the README's description completes the workflow: these metrics can be computed over different slices of data and visualized in Jupyter notebooks.

The primary language is Python, the licence is Apache 2.0, the default branch is `master`, and the last push was on 2026-08-14. The recent releases are v0.52.0 from 2026-06-12, v0.48.0 from 2025-06-23 and v0.47.1 from 2024-11-27, and the repository is not archived.

One line at the top of the README is the version caveat that should shape any adoption plan. TFMA may introduce backwards incompatible changes before version 1.0. The project is pre-1.0, and the versioning in the release list has jumped more than once, so the API surface is still moving.

Slicing: the reason to use this rather than a notebook of your own

The slicing capability is the substance of the tool, and the README's screenshot of the slicing metrics browser is the clearest visual description of it.

A slice is a subset of the evaluation data defined by a feature, and the analysis produces a metric for every slice rather than one metric for the whole set. The workflow that follows is then straightforward: define the slices you care about, run the evaluation over the data, and read a table or a chart that shows the metric for each slice side by side with the overall metric.

That is a simple idea, and doing it well is not simple, which is why the library exists. A slice over a categorical feature is easy; a slice over a numeric feature means binning, and the choice of bins changes the answer; a slice over a feature that is only available at training time and not at serving time is a modelling problem rather than an analysis one. TFMA gives you the machinery, and the design decisions about which slices to define are still yours.

The two other properties from the README follow from the same design. Because the metrics come from the trainer, the evaluation is consistent with training by construction. And because the work is distributed, the same code path runs on a laptop over a sample and on a cluster over a warehouse, which is what makes a production evaluation reproducible from a development environment.

The visualisation story is the other half of the value. The metrics browser is a Jupyter widget, so the analysis lives next to the code that produced it, and the results are interactive rather than a static chart pasted into a document. The repository includes a `notebook` directory structure and the setup.py in the tree exists specifically to package the widget, based on the template generated from the jupyter-widget cookiecutter, which tells you the notebook integration is a first class part of the project rather than an add-on.

Installing from PyPI, from nightly, or from a git ref

The recommended install is the PyPI package, and the README gives the full set of alternatives, which is useful because the choice depends on whether you need a specific commit or the latest metric behaviour.

pip install tensorflow-model-analysis

For the nightly builds, which come from the TensorFlow nightly index:

pip install -i https://pypi-nightly.tensorflow.org/simple tensorflow-model-analysis

From the head of the git repository:

pip install git+https://github.com/tensorflow/model-analysis.git#egg=tensorflow_model_analysis

And from a released version directly from git, with the tag given explicitly:

pip install git+https://github.com/tensorflow/[email protected]#egg=tensorflow_model_analysis

For local development, where you have cloned the repository and want to test your own change:

pip install -e $FOLDER_OF_THE_LOCAL_LOCATION

Two constraints are stated for the source installs. Protobuf must be installed correctly, because building TFMA from source requires protoc and all of its includes reference-able, with a link to the protobuf install instructions. And TFMA requires that TensorFlow is installed but does not have an explicit dependency on the TensorFlow PyPI package, with a pointer to the TensorFlow install guides.

That second point is worth internalising. The absence of an explicit dependency is convenient, because it lets TFMA work against a TensorFlow version your project pins, and it is also a source of installation puzzles, because pip will not tell you that TensorFlow is missing. In a clean environment, install TensorFlow first and TFMA second.

The Python version is constrained as well. The badge in the README lists Python 3.9, 3.10 and 3.11, which is narrower than the current Python ecosystem, and it is one more reason to check the environment before you upgrade anything else.

Building from source and running the tests

The build-from-source instructions exist for the case where you need a metric that is not released, and they are explicit about the tooling.

First, install protoc as the link instructs. Then create a virtual environment:

python3 -m venv <virtualenv_name>
source <virtualenv_name>/bin/activate
pip3 install setuptools wheel
git clone https://github.com/tensorflow/model-analysis.git
cd model-analysis
python3 setup.py bdist_wheel

That produces the wheel in the `dist` directory, and the install step is:

cd dist
pip3 install tensorflow_model_analysis-<version>-py3-none-any.whl

The `setup.py` in the tree shows what the build actually does, and it is more than a packaging script. It looks for the Protocol Compiler in a sequence of places, starting with the `PROTOC` environment variable and then a series of relative paths, and it customises the build to produce the Jupyter notebook integration, since the widget is based on the template generated from the jupyter-widget cookiecutter. The protobuf and protoc machinery is not incidental to the packaging, it is the reason the packaging is custom.

The test suite is run with the standard library runner, from the root project directory:

python -m unittest discover -p *_test.py

For a project this size, that is a lean test setup, and the pattern tells you something about how the code is organised. The `third_party/` directory in the tree, the `WORKSPACE` file for Bazel, the `pyproject.toml` with an extensive ruff rule selection covering pycodestyle, Pyflakes, pyupgrade, bugbear, simplify, isort, pep8 naming, pydocstyle, annotations, flake8-pytest, flake8-return, flake8-unused-arguments, flake8-fixme, flake8-eradicate, pandas-vet and numpy-specific rules, and a `pre-commit-config.yaml` at the root.

The `RELEASE.md` file and the `mkdocs.yml` with a `docs/` directory and `requirements-docs.txt` round out a package that takes release engineering and documentation seriously for a pre-1.0 project.

The notebook extension and the JupyterLab version matrix

The Jupyter integration is where most of the installation friction lives, and the README spends most of its length on it.

The first warning is about pip itself. Because of a pip issue, `pip install` might never finish, and in that case you should revert pip to version 19 instead of 20 with `pip install "pip<20"`. That is an old piece of advice attached to an old pip bug, and it is the kind of thing that will confuse a reader who arrives today.

Beyond that, using a JupyterLab extension requires installing dependencies on the command line, which can be done in the console in the JupyterLab UI or on the command line. That includes separately installing any pip package dependencies and JupyterLab labextension plugin dependencies, and the version numbers must be compatible. The examples use 0.32.0.

For JupyterLab 3.0.x:

pip install tensorflow_model_analysis==0.32.0
jupyter labextension install [email protected]
pip install jupyterlab_widgets==1.0.0

For JupyterLab 2.2.x the same two lines apply, plus `jupyter labextension install @jupyter-widgets/jupyterlab-manager@2`, and for JupyterLab 1.2.x the manager is at `@1.1`.

For the classic Jupyter Notebook, whether started with `jupyter notebook` or through the JupyterLab UI, you also need:

jupyter nbextension enable --py widgetsnbextension
jupyter nbextension enable --py tensorflow_model_analysis

with a note that if Jupyter notebook is already installed in your home directory you add `--user` to these commands, and if Jupyter is installed as root, or using a virtual environment, the parameter `--sys-prefix` might be required.

One extra requirement appears when you build from source and want to use the UI in JupyterLab: the source has to contain valid version numbers, and the Python package version number and npm package version number have to be exactly the same, and both have to be valid version numbers, for example by removing the `-dev` suffix.

The troubleshooting section offers two commands, `pip list` to check the pip packages, and `jupyter labextension list` for JupyterLab or `jupyter nbextension list` for the classic notebook, which is the fastest way to diagnose a widget that does not appear.

embed_minimal_html and the Kubeflow dependency on a CDN

Two integrations in the README show how TFMA is used beyond a personal notebook.

The first is a standalone HTML page. The TFMA notebook extension can be built into a standalone HTML file that also bundles the data into the HTML file, using the `embed_minimal_html` function documented in the Jupyter Widgets docs. The practical effect is that an evaluation report becomes a single file you can attach to a ticket, email to a stakeholder, or archive, with the data inside it rather than next to it.

The second is Kubeflow Pipelines, which includes integrations that embed the TFMA notebook extension, and the README is explicit about a dependency that catches people out. This integration relies on network access at runtime to load a variant of the JavaScript build published on unpkg.com, with the config and loader code both linked from the repository.

That is a genuine architectural constraint rather than a passing remark. A pipeline that runs in an air-gapped cluster, or one whose outputs are meant to be reproducible without a network, will not render the visualisation as written, and the fix is to vendor the JavaScript yourself. The README points at the webpack config and the widget loader so you can see exactly which file is fetched and adjust the build.

The `docs/` directory, the `mkdocs.yml` and the `_pkgdown.yml` style tooling in the tree show that the documentation is a published site rather than a directory of markdown files, and the top level link in the README points at the TensorFlow documentation site page for model analysis, which is where the API reference lives.

Editorial conclusion

Adopt TFMA if you are a TensorFlow team that needs to evaluate a model on a large slice of production data and reason about its performance per segment, because reusing the trainer's metric definitions is what keeps the evaluation honest, and the slicing model is the part that a hand written evaluation script does badly. Do not adopt it for a non-TensorFlow stack, for a one-off evaluation on a small dataset, or in a pipeline that needs a stable API, because the project explicitly warns that backwards incompatible changes can land before 1.0. Verify three things first: that the metrics you want to slice on can be produced by the model at evaluation time, that the Python and TensorFlow versions in your environment match what the installed TFMA expects, and that the notebook extension and its labextension are installed at matching version numbers, since the README requires the Python package version and the npm package version to be exactly the same.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. Releases
5. tensorflow/model-analysis on GitHub