TensorFlow Cheat Sheet
1. Introduction
TensorFlow is Google's open-source end-to-end machine learning platform. Keras serves as its high-level API, offering concise interfaces for model building, training, and deployment. TensorFlow 2.x enables Eager Execution by default, allowing intuitive tensor operations similar to NumPy.
import tensorflow as tf
print(tf.__version__) # e.g. 2.16.1
print(tf.executing_eagerly()) # True
2. Installation
# CPU only
pip install tensorflow
# GPU support (CUDA required)
pip install tensorflow[and-cuda]
# Verify GPU
python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
GPU version requires NVIDIA CUDA Toolkit and cuDNN. Using conda or Docker for environment management is recommended.
3. Tensors
Creating Tensors
# Constants
a = tf.constant([[1, 2], [3, 4]]) # shape (2,2), dtype int32
b = tf.constant([1.0, 2.0], dtype=tf.float32)
# Zeros / Ones
z = tf.zeros((3, 4)) # 3x4 of 0.0
o = tf.ones((2, 3)) # 2x3 of 1.0
# Random
r = tf.random.normal((3, 3), mean=0, stddev=1)
u = tf.random.uniform((2, 2), minval=0, maxval=10)
# Range
seq = tf.range(0, 10, delta=2) # [0, 2, 4, 6, 8]
# From NumPy
import numpy as np
t = tf.constant(np.array([1, 2, 3])) # NumPy -> Tensor
n = t.numpy() # Tensor -> NumPy
Tensors are immutable multi-dimensional arrays. Use tf.Variable for mutable tensors (model parameters).
Shape Operations
x = tf.constant([[1,2,3],[4,5,6]]) # shape (2,3)
tf.reshape(x, (3, 2)) # reshape to (3,2)
tf.reshape(x, (-1,)) # flatten to (6,)
tf.expand_dims(x, axis=0) # (1,2,3)
tf.squeeze(tf.zeros((1,3,1))) # (3,) remove size-1 dims
# Concat & Stack
a = tf.constant([[1,2]])
b = tf.constant([[3,4]])
tf.concat([a, b], axis=0) # (2,2) along rows
tf.stack([a, b], axis=0) # (2,1,2) new dim
Math Operations
a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
tf.matmul(a, b) # matrix multiply
tf.reduce_sum(a) # 10.0 (sum all)
tf.reduce_mean(a, axis=1) # [1.5, 3.5] per row
tf.reduce_max(a, axis=0) # [3.0, 4.0] per col
tf.math.log(a) # element-wise log
tf.nn.softmax(a, axis=1) # softmax per row
GPU Device
# List available devices
gpus = tf.config.list_physical_devices('GPU')
print(f"GPUs available: {len(gpus)}")
# Place ops on a specific device
with tf.device('/GPU:0'):
x = tf.random.normal((1000, 1000))
y = tf.matmul(x, x)
# Memory growth (prevent OOM)
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
TensorFlow automatically uses available GPUs by default. Use tf.device to explicitly place operations.
GradientTape
x = tf.Variable(3.0)
with tf.GradientTape() as tape:
y = x ** 2 + 2 * x + 1 # y = x^2 + 2x + 1
dy_dx = tape.gradient(y, x) # dy/dx = 2x + 2 = 8.0
print(dy_dx) # tf.Tensor(8.0, ...)
GradientTape records forward computation for automatic differentiation. It is the core of custom training loops.
4. Keras Model Building
Sequential API
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dropout(0.3),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()
Sequential is ideal for simple linear stacking. Layers are added in order; output from one flows directly into the next.
Functional API
# Multi-input model
input_text = tf.keras.Input(shape=(100,), name='text_input')
input_meta = tf.keras.Input(shape=(5,), name='meta_input')
x = tf.keras.layers.Dense(64, activation='relu')(input_text)
x = tf.keras.layers.Dropout(0.3)(x)
y = tf.keras.layers.Dense(16, activation='relu')(input_meta)
combined = tf.keras.layers.Concatenate()([x, y])
output = tf.keras.layers.Dense(1, activation='sigmoid')(combined)
model = tf.keras.Model(inputs=[input_text, input_meta], outputs=output)
Functional API supports multi-input/output, shared layers, and non-linear topologies (e.g., residual connections).
Model Subclassing
class MyModel(tf.keras.Model):
def __init__(self):
super().__init__()
self.dense1 = tf.keras.layers.Dense(128, activation='relu')
self.dropout = tf.keras.layers.Dropout(0.3)
self.dense2 = tf.keras.layers.Dense(10)
def call(self, inputs, training=False):
x = self.dense1(inputs)
x = self.dropout(x, training=training)
return self.dense2(x)
model = MyModel()
Subclassing offers maximum flexibility for research and custom forward logic. The training argument controls Dropout/BN behavior.
Common Layers
| Layer | Code | Use Case |
|---|---|---|
| Dense | Dense(64, activation='relu') | Fully connected |
| Conv2D | Conv2D(32, (3,3), activation='relu') | Image convolution |
| MaxPooling2D | MaxPooling2D(pool_size=(2,2)) | Downsampling |
| LSTM | LSTM(64, return_sequences=True) | Sequence modeling |
| Embedding | Embedding(vocab_size, 128) | Word embeddings |
| Dropout | Dropout(0.3) | Regularization |
| BatchNormalization | BatchNormalization() | Normalize activations |
| Flatten | Flatten() | Flatten multi-dim tensor |
| GlobalAveragePooling2D | GlobalAveragePooling2D() | Global avg pooling |
5. Training
model.compile()
model.compile(
optimizer='adam', # or tf.keras.optimizers.Adam(1e-3)
loss='sparse_categorical_crossentropy', # for integer labels
metrics=['accuracy']
)
model.fit()
history = model.fit(
x_train, y_train,
epochs=20,
batch_size=32,
validation_split=0.2,
callbacks=[
tf.keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
tf.keras.callbacks.ModelCheckpoint('best_model.keras', save_best_only=True),
]
)
# Evaluate
loss, acc = model.evaluate(x_test, y_test)
print(f"Test accuracy: {acc:.4f}")
# Predict
predictions = model.predict(x_new) # returns NumPy array
Custom Training Loop
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
@tf.function # compile to graph for speed
def train_step(x, y):
with tf.GradientTape() as tape:
logits = model(x, training=True)
loss = loss_fn(y, logits)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss
for epoch in range(10):
for x_batch, y_batch in train_dataset:
loss = train_step(x_batch, y_batch)
print(f"Epoch {epoch+1}, Loss: {loss:.4f}")
The @tf.function decorator compiles eager code into a computation graph for significant performance gains.
Callbacks
| Callback | Code | Description |
|---|---|---|
| EarlyStopping | EarlyStopping(patience=5, monitor='val_loss') | Stop when val_loss stops improving |
| ModelCheckpoint | ModelCheckpoint('best.keras', save_best_only=True) | Save best model weights |
| TensorBoard | TensorBoard(log_dir='./logs') | Visualize training metrics |
| ReduceLROnPlateau | ReduceLROnPlateau(factor=0.5, patience=3) | Auto reduce learning rate |
| LearningRateScheduler | LearningRateScheduler(lambda e: 1e-3 * 0.9**e) | Custom LR schedule |
6. Common Patterns
Image Classification (CNN)
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(224,224,3)),
tf.keras.layers.MaxPooling2D((2,2)),
tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
tf.keras.layers.MaxPooling2D((2,2)),
tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(10, activation='softmax')
])
# Data augmentation
data_augmentation = tf.keras.Sequential([
tf.keras.layers.RandomFlip('horizontal'),
tf.keras.layers.RandomRotation(0.1),
tf.keras.layers.RandomZoom(0.1),
])
Text Classification (Embedding + LSTM)
model = tf.keras.Sequential([
tf.keras.layers.Embedding(vocab_size, 128, input_length=max_len),
tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
Transfer Learning
# Use pre-trained MobileNetV2
base_model = tf.keras.applications.MobileNetV2(
input_shape=(224, 224, 3),
include_top=False, # remove classification head
weights='imagenet'
)
base_model.trainable = False # freeze base
model = tf.keras.Sequential([
base_model,
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(num_classes, activation='softmax')
])
# Fine-tuning: unfreeze last N layers
base_model.trainable = True
for layer in base_model.layers[:-20]:
layer.trainable = False
Available models: ResNet50, VGG16, EfficientNet, InceptionV3, etc., all under tf.keras.applications.
tf.data.Dataset
# From tensors
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.shuffle(1000).batch(32).prefetch(tf.data.AUTOTUNE)
# From directory (images)
dataset = tf.keras.utils.image_dataset_from_directory(
'data/train/',
image_size=(224, 224),
batch_size=32,
label_mode='categorical'
)
# From CSV
dataset = tf.data.experimental.make_csv_dataset(
'data.csv', batch_size=32, label_name='target'
)
tf.data provides efficient data pipelines. prefetch(AUTOTUNE) automatically optimizes CPU/GPU parallelism.
Model Saving & Loading
# SavedModel format (recommended)
model.save('my_model') # directory
loaded = tf.keras.models.load_model('my_model')
# Keras format (.keras)
model.save('model.keras')
loaded = tf.keras.models.load_model('model.keras')
# Weights only
model.save_weights('weights.h5')
model.load_weights('weights.h5')
# Export for TF Serving
tf.saved_model.save(model, 'export/1/')
7. Loss Functions
| Task | Loss Function | Code |
|---|---|---|
| Binary Classification | Binary Crossentropy | loss='binary_crossentropy' |
| Multi-class (int labels) | Sparse Categorical CE | loss='sparse_categorical_crossentropy' |
| Multi-class (one-hot) | Categorical CE | loss='categorical_crossentropy' |
| Regression | MSE | loss='mse' |
| Regression (robust) | Huber | loss=tf.keras.losses.Huber(delta=1.0) |
| Multi-label | Binary CE (sigmoid) | loss='binary_crossentropy' + sigmoid |
| Contrastive | Cosine Similarity | loss=tf.keras.losses.CosineSimilarity() |
8. Optimizers
| Optimizer | Code | Best For |
|---|---|---|
| Adam | tf.keras.optimizers.Adam(learning_rate=1e-3) | General default choice |
| SGD | tf.keras.optimizers.SGD(lr=0.01, momentum=0.9) | CV, fine-tuning |
| RMSprop | tf.keras.optimizers.RMSprop(lr=1e-3) | RNN / sequence models |
| AdamW | tf.keras.optimizers.AdamW(lr=1e-3, weight_decay=0.01) | Transformers / large models |
| Adagrad | tf.keras.optimizers.Adagrad(lr=0.01) | Sparse features / NLP |
9. TensorFlow vs PyTorch
| Feature | TensorFlow / Keras | PyTorch |
|---|---|---|
| Execution Mode | Eager + @tf.function graph compilation | Eager + torch.compile |
| High-level API | tf.keras (built-in) | Lightning / HuggingFace |
| Deployment | TF Serving, TF Lite, TF.js | TorchServe, ONNX, ExecuTorch |
| Mobile | TF Lite (mature) | ExecuTorch (newer) |
| Browser | TensorFlow.js | ONNX.js / Transformers.js |
| Community Trend | Preferred for production | Preferred for research |
| Debugging | Eager mode + TF Debugger | Native Python debugging |
| Distributed Training | tf.distribute.Strategy | torch.distributed |
Both frameworks are mature. TensorFlow has a more complete deployment ecosystem, while PyTorch is more popular in research.
10. Deployment
TF Serving
# Save model in SavedModel format
tf.saved_model.save(model, 'export/model/1/')
# Docker: run TF Serving
# docker pull tensorflow/serving
# docker run -p 8501:8501 \
# --mount type=bind,source=$(pwd)/export/model,target=/models/my_model \
# -e MODEL_NAME=my_model tensorflow/serving
# REST API prediction
import requests, json
data = json.dumps({"instances": x_test[:3].tolist()})
resp = requests.post('http://localhost:8501/v1/models/my_model:predict',
data=data, headers={"Content-Type": "application/json"})
print(resp.json()['predictions'])
TF Lite (Mobile)
# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_saved_model('export/model/1/')
converter.optimizations = [tf.lite.Optimize.DEFAULT] # quantize
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
# Inference with TFLite
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output = interpreter.get_tensor(interpreter.get_output_details()[0]['index'])
TensorFlow.js (Browser)
# Convert to TF.js format
# pip install tensorflowjs
# tensorflowjs_converter --input_format=tf_saved_model \
# export/model/1/ tfjs_model/
# In browser JavaScript:
# const model = await tf.loadGraphModel('tfjs_model/model.json');
# const input = tf.tensor2d([[...features...]]);
# const prediction = model.predict(input);
# prediction.print();
11. Related Tools
12. FAQ
What are the main differences between TensorFlow 2.x and 1.x?
TF 2.x uses Eager Execution by default, removes tf.Session, adopts Keras as the official high-level API, and provides graph compilation via @tf.function. The API is more concise and Pythonic.
When to use Sequential vs Functional vs Subclassing?
Sequential for simple linear models; Functional API for complex models with multiple inputs/outputs and shared layers; Subclassing for research scenarios requiring custom forward logic (e.g., GANs, custom attention).
How to fix GPU out-of-memory (OOM) errors?
1) Enable memory growth: tf.config.experimental.set_memory_growth(gpu, True); 2) Reduce batch_size; 3) Use mixed precision training tf.keras.mixed_precision; 4) Use gradient accumulation.
Is TensorFlow good for production deployment?
Yes, TensorFlow has the most comprehensive deployment ecosystem: TF Serving (server), TF Lite (mobile/embedded), TF.js (browser). Google internally uses TensorFlow at scale serving billions of requests.
How to choose between TensorFlow and PyTorch?
Choose TensorFlow for deployment convenience, mobile support, or browser execution. Choose PyTorch for academic research, rapid prototyping, or more flexible debugging. Both handle most deep learning tasks well.