TensorFlow

import tensorflow as tf
from tensorflow import keras

# Sequential model for simple architectures
model = keras.Sequential([
    keras.layers.Input(shape=(784,)),
    keras.layers.Dense(256, activation="relu"),
    keras.layers.Dropout(0.3),
    keras.layers.Dense(128, activation="relu"),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10, activation="softmax"),
])

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=1e-3),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

# Train
history = model.fit(
    x_train, y_train,
    epochs=20,
    batch_size=64,
    validation_split=0.2,
    callbacks=[
        keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
        keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=2),
        keras.callbacks.ModelCheckpoint("best_model.keras", save_best_only=True),
    ],
)

# Multi-input, multi-output model
text_input = keras.Input(shape=(None,), dtype="int32", name="text")
image_input = keras.Input(shape=(224, 224, 3), name="image")

# Text branch
x = keras.layers.Embedding(vocab_size, 128)(text_input)
x = keras.layers.LSTM(64)(x)

# Image branch
y = keras.applications.EfficientNetV2B0(include_top=False, pooling="avg")(image_input)
y = keras.layers.Dense(128, activation="relu")(y)

# Combine
combined = keras.layers.Concatenate()([x, y])
combined = keras.layers.Dense(64, activation="relu")(combined)

# Multiple outputs
category = keras.layers.Dense(num_categories, activation="softmax", name="category")(combined)
sentiment = keras.layers.Dense(1, activation="sigmoid", name="sentiment")(combined)

model = keras.Model(
    inputs=[text_input, image_input],
    outputs=[category, sentiment],
)

# Fine-grained control over training
@tf.function                              # traces the Python step into a TensorFlow graph
def train_step(model, optimizer, x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = loss_fn(y, predictions)

    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# Training loop
for epoch in range(num_epochs):
    for batch_x, batch_y in train_dataset:
        loss = train_step(model, optimizer, batch_x, batch_y)
    
    # Validation
    val_loss = tf.reduce_mean([
        loss_fn(y, model(x, training=False))
        for x, y in val_dataset
    ])
    print(f"Epoch {epoch}: loss={loss:.4f}, val_loss={val_loss:.4f}")

# Save model
model.save("my_model.keras")              # Keras format
model.export("models/my_model/1")        # SavedModel for TF Serving (model.save("dir/") raises ValueError in Keras 3)

# TFLite for mobile
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # dynamic-range quantization
# Google is moving on-device inference to LiteRT (pip package ai-edge-litert); tf.lite.TFLiteConverter still works in 2.21
tflite_model = converter.convert()
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

# TensorFlow Serving (Docker): the version number (1) must be a directory under the model name
# docker run -p 127.0.0.1:8501:8501 --mount type=bind,source="$PWD/models",target=/models \
#   -e MODEL_NAME=my_model tensorflow/serving

# REST API inference
import requests
response = requests.post(
    "http://localhost:8501/v1/models/my_model:predict",
    json={"instances": x_test[:5].tolist()},
)
predictions = response.json()["predictions"]

python3 -m venv .venv && source .venv/bin/activate
pip install tensorflow                     # CPU only (Python 3.10+)
pip install "tensorflow[and-cuda]"         # Linux NVIDIA GPU: pulls CUDA 12 and cuDNN 9 wheels
pip install tensorboard                    # no longer installed with tensorflow 2.21
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"