tensorflow-guide

TensorFlow Guide

Overview

TensorFlow is a production-grade machine learning framework that excels at deployment, distributed training, and hardware acceleration. While PyTorch dominates pure research prototyping, TensorFlow remains the standard in industry ML systems and is heavily used in applied research where models must move from experiment to production.

TensorFlow 2.x unified eager execution with graph-mode performance through tf.function, but this hybrid approach introduces subtle pitfalls. Understanding when and how TensorFlow traces functions, manages GPU memory, and distributes computation is essential for writing correct and efficient code.

This guide covers the key patterns that trip up researchers: tf.function tracing semantics, GPU memory management, distributed strategies, model export, and the ecosystem of tools (TFX, TensorBoard, TF Serving) that make TensorFlow uniquely powerful for end-to-end ML workflows.

tf.function: The Critical Abstraction

How Tracing Works

import tensorflow as tf

@tf.function
def add(a, b):
    print("Tracing!")  # Runs only during tracing, NOT every call
    tf.print("Executing!")  # Runs every call (TF op)
    return a + b

# First call with float32 shape (2,) -- traces
add(tf.constant([1.0, 2.0]), tf.constant([3.0, 4.0]))  # Prints "Tracing!" + "Executing!"

# Second call with same signature -- reuses trace
add(tf.constant([5.0, 6.0]), tf.constant([7.0, 8.0]))  # Prints only "Executing!"

# Third call with different dtype -- re-traces!
add(tf.constant([1, 2]), tf.constant([3, 4]))  # Prints "Tracing!" + "Executing!"

Common tf.function Pitfalls

# PITFALL 1: Python side effects in tf.function
counter = 0

@tf.function
def increment():
    global counter
    counter += 1  # Only runs during tracing! counter stays at 1 forever.
    return counter

# FIX: Use tf.Variable for mutable state
counter = tf.Variable(0)

@tf.function
def increment():
    counter.assign_add(1)
    return counter

# PITFALL 2: Creating variables inside tf.function
@tf.function
def bad_function(x):
    w = tf.Variable(tf.random.normal([3, 3]))  # ERROR on second call!
    return x @ w

# FIX: Create variables outside, pass as arguments or use Keras layers
w = tf.Variable(tf.random.normal([3, 3]))

@tf.function
def good_function(x):
    return x @ w

# PITFALL 3: Python lists that grow
@tf.function
def bad_accumulate(dataset):
    results = []
    for x in dataset:
        results.append(x * 2)  # Creates new trace on every iteration!
    return results

# FIX: Use tf.TensorArray
@tf.function
def good_accumulate(dataset):
    results = tf.TensorArray(tf.float32, size=0, dynamic_size=True)
    for i, x in enumerate(dataset):
        results = results.write(i, x * 2)
    return results.stack()