PyTorch & TensorFlow Cheat Sheet

// PyTorch
import torch

# Initialization
t = torch.tensor([1, 2, 3])
t = torch.zeros(3, 4)
t = torch.ones(3, 4)
t = torch.randn(3, 4)  # Normal distribution
t = torch.rand(3, 4)   # Uniform [0,1)
t = torch.arange(10)
t = torch.linspace(0, 1, 10)

# Device
t = torch.tensor([1, 2], device='cuda')
t = t.to('cpu')
t = t.cuda()

// TensorFlow
import tensorflow as tf

# Initialization
t = tf.constant([1, 2, 3])
t = tf.zeros([3, 4])
t = tf.ones([3, 4])
t = tf.random.normal([3, 4])
t = tf.random.uniform([3, 4])
t = tf.range(10)
t = tf.linspace(0.0, 1.0, 10)

# Device
with tf.device('/GPU:0'):
    t = tf.constant([1, 2])

// PyTorch
t1 + t2              # Element-wise add
t1 * t2              # Element-wise multiply
torch.matmul(t1, t2) # Matrix multiply
t.T                  # Transpose
t.reshape(3, 4)      # Reshape
t.view(3, 4)         # View (similar to reshape)
t.squeeze()          # Remove dims of size 1
t.unsqueeze(0)       # Add dim of size 1
t.cat([t1, t2], dim=0)  # Concatenate
t.stack([t1, t2])    # Stack into new dimension

// TensorFlow
t1 + t2              # Element-wise add
t1 * t2              # Element-wise multiply
tf.matmul(t1, t2)    # Matrix multiply
tf.transpose(t)      # Transpose
tf.reshape(t, [3, 4]) # Reshape
tf.squeeze(t)        # Remove dims of size 1
tf.expand_dims(t, 0) # Add dim of size 1
tf.concat([t1, t2], axis=0)  # Concatenate
tf.stack([t1, t2])   # Stack into new dimension

// Both
# Element-wise operations
sqrt, exp, log, sin, cos, abs, sign, etc.

// PyTorch (Class-based)
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(32*14*14, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)  # Flatten
        x = self.fc(x)
        return x

model = CNN()

// TensorFlow (Functional API)
import tensorflow as tf
from tensorflow import keras

model = keras.Sequential([
    keras.layers.Conv2D(32, (3,3), padding='same', activation='relu', input_shape=(28,28,3)),
    keras.layers.MaxPooling2D((2,2)),
    keras.layers.Flatten(),
    keras.layers.Dense(10, activation='softmax')
])

// TensorFlow (Subclassing, like PyTorch)
class CNN(keras.Model):
    def __init__(self):
        super().__init__()
        self.conv1 = keras.layers.Conv2D(32, (3,3), padding='same')
        self.pool = keras.layers.MaxPooling2D((2,2))
        self.fc = keras.layers.Dense(10)

    def call(self, x, training=False):
        x = self.conv1(x)
        x = tf.nn.relu(x)
        x = self.pool(x)
        x = tf.keras.layers.Flatten()(x)
        return self.fc(x)

// PyTorch (Manual loop)
import torch.optim as optim

model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    for X_batch, y_batch in train_loader:
        # Forward pass
        y_pred = model(X_batch)
        loss = criterion(y_pred, y_batch)

        # Backward pass
        optimizer.zero_grad()
        loss.backward()

        # Update
        optimizer.step()

    # Validation
    with torch.no_grad():
        val_loss = 0
        for X_val, y_val in val_loader:
            y_pred = model(X_val)
            val_loss += criterion(y_pred, y_val)

// TensorFlow (Using fit)
model = CNN()
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

history = model.fit(
    train_dataset,
    epochs=10,
    validation_data=val_dataset,
    callbacks=[keras.callbacks.EarlyStopping(patience=3)]
)

// TensorFlow (Custom training loop)
@tf.function
def train_step(X, y):
    with tf.GradientTape() as tape:
        y_pred = model(X, training=True)
        loss = criterion(y, y_pred)

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

for epoch in range(10):
    for X_batch, y_batch in train_dataset:
        train_step(X_batch, y_batch)

// PyTorch
from torch.utils.data import Dataset, DataLoader

class CustomDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

dataset = CustomDataset(X_train, y_train)
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

for X_batch, y_batch in train_loader:
    # Train on batch
    pass

// TensorFlow
# From numpy
train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=1000)
train_dataset = train_dataset.batch(32)
train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)

for X_batch, y_batch in train_dataset:
    # Train on batch
    pass

// Image Data
from tensorflow.keras.preprocessing.image import ImageDataGenerator

datagen = ImageDataGenerator(
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True
)

train_dataset = datagen.flow_from_directory(
    'data/train',
    target_size=(224, 224),
    batch_size=32
)

// PyTorch Layers
nn.Linear(in, out)
nn.Conv2d(in_channels, out_channels, kernel_size, padding)
nn.ReLU()
nn.Sigmoid()
nn.Tanh()
nn.Dropout(p=0.5)
nn.BatchNorm2d(num_features)
nn.LayerNorm(normalized_shape)
nn.Embedding(num_embeddings, embedding_dim)
nn.LSTM(input_size, hidden_size, num_layers)

// PyTorch Losses
nn.CrossEntropyLoss()  # Classification
nn.BCELoss()           # Binary classification
nn.MSELoss()           # Regression
nn.L1Loss()            # L1 regression
nn.SmoothL1Loss()      # Huber

// TensorFlow Losses
keras.losses.SparseCategoricalCrossentropy()
keras.losses.BinaryCrossentropy()
keras.losses.MeanSquaredError()
keras.losses.MeanAbsoluteError()
keras.losses.Huber()

// PyTorch Optimizers
optim.SGD(params, lr=0.01, momentum=0.9)
optim.Adam(params, lr=1e-3, betas=(0.9, 0.999))
optim.AdamW(params, lr=1e-3, weight_decay=1e-5)
optim.RMSprop(params, lr=0.01)

// TensorFlow Optimizers
keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
keras.optimizers.Adam(learning_rate=1e-3)
keras.optimizers.RMSprop(learning_rate=0.01)

// PyTorch - Save/Load
torch.save(model.state_dict(), 'model.pth')
model.load_state_dict(torch.load('model.pth'))

// TensorFlow - Save/Load
model.save('model.h5')
model.save('model/')  # SavedModel format
model = keras.models.load_model('model.h5')

// PyTorch - Evaluation
model.eval()
with torch.no_grad():
    y_pred = model(X_test)

// TensorFlow - Evaluation
y_pred = model.predict(X_test)
loss, accuracy = model.evaluate(X_test, y_test)

// Metrics (both)
Accuracy: Correct predictions / total
Precision, Recall, F1
Confusion matrix
ROC curve, AUC

// Common patterns
model.train()      # Training mode (enables dropout, BN)
model.eval()       # Evaluation mode
requires_grad=False  # Freeze parameters
detach()           # Remove from computation graph (PyTorch)
.numpy()           # Convert to numpy (PyTorch)
.numpy()           # Convert TensorFlow tensor to numpy