Tensorflow Deep Learning

TensorFlow and Deep Learning

This skill covers project structure, model development, training, evaluation, and deployment practices for neural networks built with TensorFlow and Keras.

Workflow for Training and Shipping a Model

1. Separate concerns into modules — Keep data loading, model definition, training, evaluation, and serving code in distinct files; treat notebooks as exploratory only.
2. Build the input pipeline first — Use tf.data.Dataset for scalable, prefetching input processing; validate shapes, dtypes, label ranges, and class balance before writing any model code.
3. Split before any fitting — Create train/validation/test splits before fitting normalization statistics or augmentation parameters, to avoid leakage.
4. Start with a tiny baseline — Build a small model and run a tiny overfit test (fit on a handful of examples until loss goes to ~0) to prove the training loop works before scaling up.
5. Train with callbacks — Wire in ModelCheckpoint (save best by validation metric, not final epoch), EarlyStopping, a learning-rate schedule, and TensorBoard logging.
6. Evaluate with task-appropriate metrics — Use a held-out test set only for final reporting, never for tuning.
7. Export with a explicit signature — Save with model.export() or tf.saved_model.save() and define serving input signatures explicitly.
8. Smoke test the export — Load the exported model and run inference on sample inputs before deploying.

Project Structure

• Separate data loading, model definition, training, evaluation, and serving code into distinct modules (e.g., data.py, model.py, train.py, evaluate.py, serve.py).
• Keep model hyperparameters in typed config objects (dataclasses) or config files, not scattered as literals through training code.
• Store checkpoints, logs, and exported models outside source directories (e.g., under artifacts/ or a configured output path), so they don't get swept into version control accidentally.
• Keep notebooks exploratory; once an approach is validated, move the repeatable training code into modules that can be run as scripts and covered by tests.

project/
  data.py          # tf.data pipeline construction
  model.py         # Keras model definition
  config.py        # typed hyperparameter config
  train.py         # training loop / Keras fit orchestration
  evaluate.py       # metrics computation on held-out data
  serve.py         # export + inference smoke test
  artifacts/       # checkpoints, logs, saved models (gitignored)

Input Pipeline (tf.data)

import tensorflow as tf

def build_dataset(file_pattern: str, batch_size: int, shuffle: bool) -> tf.data.Dataset:
    files = tf.data.Dataset.list_files(file_pattern, shuffle=shuffle)
    ds = files.interleave(
        tf.data.TFRecordDataset,
        cycle_length=tf.data.AUTOTUNE,
        num_parallel_calls=tf.data.AUTOTUNE,
    )
    ds = ds.map(_parse_example, num_parallel_calls=tf.data.AUTOTUNE)
    if shuffle:
        ds = ds.shuffle(buffer_size=10_000)
    ds = ds.batch(batch_size, drop_remainder=False)
    return ds.prefetch(tf.data.AUTOTUNE)

def _parse_example(record: tf.Tensor) -> tuple[tf.Tensor, tf.Tensor]:
    feature_spec = {
        "features": tf.io.FixedLenFeature([32], tf.float32),
        "label": tf.io.FixedLenFeature([], tf.int64),
    }
    parsed = tf.io.parse_single_example(record, feature_spec)
    return parsed["features"], parsed["label"]

• Use num_parallel_calls=tf.data.AUTOTUNE on .map()/.interleave() and always end the pipeline with .prefetch(tf.data.AUTOTUNE).
• Shuffle before batching, with a buffer large enough to approximate a full shuffle for the dataset size.
• Validate data shapes, dtypes, label ranges, and class balance on a sample batch before starting a long training run.

Model Development

• Use Keras layers and models (tf.keras.Model, tf.keras.Sequential) unless a specific need requires lower-level tf.GradientTape control.
• Prefer explicit input shapes and named inputs/outputs — this makes serving signatures and debugging easier.
• Start with a small baseline model and confirm it can overfit a tiny subset of data (loss near zero) before scaling depth, width, or dataset size — this catches pipeline and loss-function bugs early.
• Pin random seeds (tf.random.set_seed, numpy.random.seed) where reproducibility matters, but document that GPU execution can still be nondeterministic for some ops (e.g., certain cuDNN reductions).
• Only enable mixed precision (tf.keras.mixed_precision.set_global_policy("mixed_float16")) after validating numerical stability at full precision — some losses (e.g., ones involving log or exp) need loss scaling to avoid underflow.

import tensorflow as tf

def build_model(input_dim: int, num_classes: int) -> tf.keras.Model:
    inputs = tf.keras.Input(shape=(input_dim,), name="features")
    x = tf.keras.layers.Dense(128, activation="relu")(inputs)
    x = tf.keras.layers.Dropout(0.3)(x)
    x = tf.keras.layers.Dense(64, activation="relu")(x)
    outputs = tf.keras.layers.Dense(num_classes, activation="softmax", name="predictions")(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs, name="baseline_classifier")

Training

import tensorflow as tf

model = build_model(input_dim=32, num_classes=10)
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

callbacks = [
    tf.keras.callbacks.ModelCheckpoint(
        filepath="artifacts/checkpoints/best.keras",
        monitor="val_accuracy",
        save_best_only=True,
    ),
    tf.keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=5, restore_best_weights=True,
    ),
    tf.keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=3, min_lr=1e-6,
    ),
    tf.keras.callbacks.TensorBoard(log_dir="artifacts/logs"),
]

train_ds = build_dataset("data/train-*.tfrecord", batch_size=64, shuffle=True)
val_ds = build_dataset("data/val-*.tfrecord", batch_size=64, shuffle=False)

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=50,
    callbacks=callbacks,
)

• Split data into train/validation/test before fitting any normalization layer or augmentation parameters — fitting on the full dataset leaks test statistics into training.
• Use the validation set for tuning (architecture, learning rate, regularization); reserve a separate test set purely for final reporting, touched once.
• Track loss curves, metrics, learning rate, and resource utilization (TensorBoard, or an experiment tracker) for every run.
• Save the checkpoint with the best validation metric (save_best_only=True), not just the weights from the final epoch — the final epoch is not necessarily the best one, especially with EarlyStopping.

Evaluation

• Report task-appropriate metrics: AUROC/F1/calibration for classification, perplexity or BLEU/ROUGE for language tasks, MAE/RMSE for regression.
• Include confusion matrices or per-class error slices for classification tasks — aggregate accuracy hides class-level failures.
• Evaluate on edge cases and known distribution shifts when data allows (e.g., a held-out slice from a different time period or source).
• Compare against non-neural baselines (logistic regression, gradient-boosted trees) when the dataset is small or tabular — a deep model that underperforms a simple baseline is a signal, not a footnote.

test_ds = build_dataset("data/test-*.tfrecord", batch_size=64, shuffle=False)
results = model.evaluate(test_ds, return_dict=True)
print(results)  # {'loss': ..., 'accuracy': ...}

Deployment

• Export models with an explicit input signature so serving infrastructure (TF Serving, Vertex AI, etc.) knows the expected input shape and dtype:

model.export("artifacts/saved_model/v1")

# Or with an explicit serving signature:
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 32], dtype=tf.float32, name="features")])
def serve_fn(features):
    return {"predictions": model(features, training=False)}

tf.saved_model.save(model, "artifacts/saved_model/v1", signatures={"serving_default": serve_fn})

• Keep preprocessing identical between training and serving — bake normalization/tokenization into the exported graph (e.g., tf.keras.layers.Normalization as part of the model) rather than duplicating logic in a separate serving-side script.
• Add a smoke test that loads the exported SavedModel and runs inference on a handful of known sample inputs, asserting outputs are in the expected range/shape, before promoting a model to production.
• Monitor latency, memory, prediction drift, and input schema changes once deployed — a model that was correct at export time can silently degrade as upstream data shifts.

Common Mistakes

• Tuning architecture before verifying labels and data quality — a data bug will outlast any architecture change.
• Leaking validation or test data through preprocessing or augmentation fit on the full dataset instead of the training split alone.
• Relying on accuracy alone for imbalanced datasets — use precision/recall, F1, or AUROC instead.
• Deploying a notebook-only model with no reproducible training script or pinned dependencies.
• Ignoring batch size, dtype, or device (CPU/GPU) differences between training and inference, which can silently change numerical results (especially under mixed precision).