+1 (415) 366-1364
DeepAdvisors

Upgrading to TensorFlow 2.21: A Field Guide

If you run TensorFlow in production, there is a good chance some of it is pinned to a release from years ago. It still works, so it never got upgraded. This guide is the process we use when a client asks us to bring a codebase onto the current TensorFlow 2.21 release line: why it is worth doing, how to find out what will break, and how to do it without changing what your models predict.

Why upgrade at all

Three reasons, in the order clients usually care about them:

  1. Security and support. Only the current release line receives patches. The Python interpreters old TensorFlow builds were compiled against are themselves end-of-life, which is the thing that eventually fails a security review.
  2. Python support. TensorFlow 2.21 supports Python 3.10 through 3.13; Python 3.9 support was dropped. If the rest of your platform has moved on, an old TensorFlow pin is the thing blocking the interpreter upgrade.
  3. Performance and ecosystem. Current builds bring newer XLA, CUDA, and oneDNN paths, and they are what Keras 3, KerasHub, and LiteRT are tested against.

The argument against is always the same: risk. The rest of this guide is about managing it.

Step 1: know what you are running

Start by writing down the real versions, not the ones in the README:

python -c "import sys, tensorflow as tf; print(sys.version); print(tf.__version__)"
pip freeze | grep -iE "tensorflow|keras|numpy|protobuf"

Then grep for the APIs that are most likely to break. This one-liner gives you a quick census of the legacy surface area:

grep -rnE "tf\.compat\.v1|tf\.contrib|tf\.estimator|tf\.Session|tf\.placeholder|tf\.get_variable|tf\.layers\." --include="*.py" . | wc -l
grep -rnE "tf\.contrib\.[a-z_]+" -o --include="*.py" . | sort | uniq -c | sort -rn

Every tf.contrib module is gone in 2.x; the second command tells you which ones you depend on so you can look up where each moved (TensorFlow Addons, TensorFlow Probability, plain tf, or nowhere).

Step 2: capture a golden set before touching anything

This is the step people skip and regret. Before you change a single line, run the existing model on a fixed set of inputs and save the outputs:

import numpy as np

inputs = np.load("golden_inputs.npy")           # a few hundred real examples
outputs = legacy_predict(inputs)                  # whatever your current path is
np.save("golden_outputs_legacy.npy", outputs)

Every later phase replays the same inputs and diffs against this file:

def assert_parity(new_outputs, atol=1e-5, rtol=1e-4):
    ref = np.load("golden_outputs_legacy.npy")
    np.testing.assert_allclose(new_outputs, ref, atol=atol, rtol=rtol)

Pick the tolerance deliberately. Float32 reductions reorder across versions and hardware, so bit-exact is the wrong target; a tolerance you can defend to the model owner is the right one.

Step 3: audit the deprecated APIs

The big categories we see in client code, roughly in order of effort:

tf.compat.v1 graphs and sessions. Code that builds a graph with placeholders and runs it in a Session still executes in 2.x under compat.v1, but it is a dead end: no new features, no Keras 3, awkward serving. The modern equivalent is eager code wrapped in tf.function:

# before
x = tf.compat.v1.placeholder(tf.float32, [None, 784])
logits = build_model(x)
with tf.compat.v1.Session() as sess:
    out = sess.run(logits, feed_dict={x: batch})

# after
model = build_model()            # a keras.Model or a plain callable

@tf.function
def predict(batch):
    return model(batch, training=False)

out = predict(batch)

Estimators. tf.estimator is deprecated and not available in current builds. Convert model_fn logic into a Keras model and the input_fn into a tf.data.Dataset. The structure maps cleanly: ModeKeys.TRAIN becomes model.fit, PREDICT becomes model.predict or a tf.function, and train_and_evaluate becomes fit with a validation dataset and callbacks.

tf.layers and tf.get_variable. Replace with keras.layers and layer-owned weights. Variable scopes and reuse flags go away entirely; Keras layers are objects, so "reuse" is just calling the same layer twice.

Checkpoints. TF1 Saver checkpoints are name-based. After converting to Keras, you will need a one-time script that loads the old checkpoint with tf.train.load_checkpoint, maps variable names to the new layers, and assigns them. Run the golden-set check immediately afterward; a silent name mismatch is the most common way to ship a randomly-initialized layer.

Keras version. TensorFlow 2.16 and later ship with Keras 3 as tf.keras by default. If your code relies on Keras 2 behaviour (for example model.save to HDF5 with custom objects, or keras.backend functions that were removed), you can pin tf-keras as a bridge, but treat that as a stepping stone, not a destination. Our Keras 3 multi-backend tutorial covers the differences that matter.

Step 4: pin, build, and test in a branch

Upgrade in a branch with a container, not on a shared machine:

FROM python:3.12-slim
RUN pip install --no-cache-dir "tensorflow==2.21.*" "numpy>=1.26,<3"
COPY requirements.lock /tmp/
RUN pip install --no-cache-dir -r /tmp/requirements.lock

Two rules we hold to:

  • One upgrade at a time. If you are several majors behind, go 1.15 to 2.x under compat.v1 first (make the tests pass), then remove compat.v1, then bump to 2.21. Collapsing the steps makes every failure ambiguous.
  • Lock the whole tree. pip freeze > requirements.lock after each passing phase. protobuf and numpy in particular have broken more TensorFlow upgrades than TensorFlow itself has.

Run the unit tests, then the golden-set parity check, after every phase.

Common breakages

The ones that account for most of our time on migration engagements:

  • Random seeds. TF2 changed how global and op-level seeds interact. If a test relies on a specific random sequence, it will fail. Fix the test, not the model.
  • tf.nn.softmax_cross_entropy_with_logits argument order and labels semantics. Easy to get silently wrong; the parity check catches it.
  • tf.data iteration. make_one_shot_iterator and friends are gone; iterate the dataset directly.
  • Custom ops built against old headers. They must be rebuilt against the target TensorFlow, and the build toolchain requirements have moved.
  • Feature columns. Removed with Estimators. Replace with Keras preprocessing layers (StringLookup, Normalization, CategoryEncoding).

Modernize or rewrite?

Ask two questions. Does anyone on the team understand the model well enough to defend its behaviour? and Would a modern pre-trained backbone or a fine-tuned open-weight model plausibly beat it? If the answers are "no" and "yes", port the data pipeline and the evaluation harness, and replace the model. A faithful port of a 2017 architecture is sometimes the wrong deliverable.

Checklist

  1. Inventory versions and legacy API usage.
  2. Capture a golden input/output set.
  3. Upgrade one step at a time in a container; lock dependencies after each passing phase.
  4. Convert Estimators and compat.v1 graphs to Keras and tf.function.
  5. Map old checkpoints to new layers; re-run parity.
  6. Update serving (SavedModel export) and CI images.
  7. Decide, per model, whether to port or replace.

If you would rather have someone who does this every month run the audit, our TensorFlow Modernization & Migration engagement is a fixed-scope version of exactly this process. Contact us to talk it through.