TensorFlow: Google's Machine Learning Framework

TensorFlow is an end-to-end platform: define a model, train it, export it, and run it on a server, in a browser, on a phone, or on a microcontroller. That last part is the honest reason most teams choose it. On modelling capability the frameworks have converged; on breadth of deployment targets TensorFlow still has the widest coverage, and if your target is a browser or an embedded device the decision is largely made for you.

Eager code, graph execution

TensorFlow runs eagerly by default, like ordinary Python. tf.function traces a function once and turns it into a graph that can be optimised — operations fused, memory reused, work distributed across GPUs or TPUs — and executed without the Python interpreter in the loop.

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

Tracing has one consequence worth internalising: the Python code runs at trace time, not at every call. A print() inside a traced function fires once; a Python counter increments once; a data-dependent if on a tensor value gets baked into the graph as whichever branch was taken during tracing. Use tf.print for logging inside a graph, tf.Variable for state, and tf.cond for tensor-dependent branching.

The second consequence is retracing. A new input shape or dtype triggers a fresh trace, and a function called with varying shapes will retrace continually, which is slower than eager execution. Fixed batch shapes, or an explicit input_signature, prevent it.

Keras at three levels

Most TensorFlow work now happens through Keras, which offers three abstraction levels, and knowing which you need saves considerable time.

Sequential — a linear stack of layers. Fine for straightforward feedforward and convolutional models, and unable to express anything with branching.

Functional — layers as callables composed into a graph. Handles multiple inputs, multiple outputs, shared layers, and skip connections. This covers the large majority of real architectures and is where you should default.

Subclassing — write a Model subclass with your own call method. Full Python flexibility for genuinely dynamic architectures, at the cost of losing some static-graph checking.

model.fit() is more capable than it first appears — callbacks give you early stopping, checkpointing, learning-rate scheduling, and TensorBoard logging without a custom loop. Write a custom training loop when you need something fit genuinely cannot express, such as multiple optimisers or an adversarial setup, not as a default.

The input pipeline decides your throughput

A common and expensive surprise: GPU utilisation sits at thirty percent and the bottleneck is data loading, not the model. tf.data exists to fix this, and three calls do most of the work.

dataset = (
    tf.data.Dataset.from_tensor_slices((paths, labels))
    .shuffle(10_000)
    .map(load_and_augment, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(64)
    .prefetch(tf.data.AUTOTUNE)
)

num_parallel_calls parallelises preprocessing across CPU cores. prefetch overlaps the preparation of batch n+1 with the training of batch n, so the accelerator never waits. cache() is worth adding after map when the preprocessed data fits in memory and the augmentation is deterministic. Order matters: shuffle before batch, prefetch last.

Match the tool to the target

This is where TensorFlow earns its complexity.

Server-side inference — SavedModel plus TensorFlow Serving, which handles versioning, batching of concurrent requests, and hot model swaps behind a stable endpoint.

Browser — TensorFlow.js runs a converted model in the client, so data never leaves the device and there is no inference cost to you. The constraint is model size over the network.

Mobile — LiteRT (formerly TensorFlow Lite) converts and quantises for phones. Post-training quantisation to int8 typically shrinks a model roughly fourfold and speeds it up substantially, usually with a small accuracy cost that you must measure rather than assume.

Microcontrollers — the same lineage runs in kilobytes of memory on embedded hardware.

Decide the target before you finish the model. A model designed without its deployment constraint in mind frequently cannot be converted, or converts into something too large or too slow for the device, and the rework is real.

Common mistakes

Python side effects inside tf.function. They run once at trace time and then never again.

Retracing on every call. Varying input shapes make graph mode slower than eager.

Blaming the model for low GPU utilisation. Check the input pipeline first; prefetch and parallel map fix most of it.

Choosing the deployment target last. Conversion constraints should shape the architecture, not surprise it.