mdc-tensorflow

tensorflow Best Practices

This document outlines our team's definitive guidelines for writing TensorFlow code. Adhering to these practices ensures maintainable, performant, and reproducible machine learning systems.

1. Code Organization and Structure

Always structure your TensorFlow projects for clarity and modularity. Separate concerns into distinct files or functions.

1.1. Model Definition: Use tf.keras

Always define models using the tf.keras API. It's the high-level, declarative standard for TensorFlow. Prefer the Functional API for complex models and Sequential for simple stacks.

❌ BAD: Raw tf.Variable and custom loops for model architecture.

import tensorflow as tf

class BadCustomModel:
    def __init__(self):
        self.w1 = tf.Variable(tf.random.normal([784, 128]))
        self.b1 = tf.Variable(tf.zeros([128]))
        # ... more manual variables

✅ GOOD: tf.keras.Model or tf.keras.Sequential.

import tensorflow as tf

def build_model(input_shape: tuple[int, ...], num_classes: int) -> tf.keras.Model:
    inputs = tf.keras.Input(shape=input_shape)
    x = tf.keras.layers.Flatten()(inputs)
    x = tf.keras.layers.Dense(128, activation='relu')(x)
    x = tf.keras.layers.Dropout(0.2)(x)
    outputs = tf.keras.layers.Dense(num_classes, activation='softmax')(x)
    return tf.keras.Model(inputs=inputs, outputs=outputs)

# Usage:
model = build_model((28, 28), 10)

1.2. Data Pipelines: Use tf.data

Always use the tf.data API for building robust and efficient input pipelines. This is critical for scaling to large datasets and optimizing I/O.

❌ BAD: Loading all data into memory or using numpy arrays for large datasets.

import numpy as np
# ... load huge_data into numpy array
# x_train, y_train = np.load('huge_data.npy')
# model.fit(x_train, y_train, batch_size=32) # Inefficient for large data

✅ GOOD: Stream data with tf.data.Dataset.

import tensorflow as tf
import tensorflow_datasets as tfds

def create_dataset(split: str, batch_size: int) -> tf.data.Dataset:
    ds = tfds.load('mnist', split=split, as_supervised=True)
    ds = ds.map(lambda img, label: (tf.cast(img, tf.float32) / 255.0, label))
    ds = ds.shuffle(1024).batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return ds

# Usage:
train_ds = create_dataset('train', 128)
test_ds = create_dataset('test', 128)