AI & ML Development Cheatsheet

import torch
import torch.nn as nn

# Tensor creation
x = torch.tensor([1, 2, 3], dtype=torch.float32)
z = torch.zeros(3, 4)
r = torch.randn(3, 4)           # standard normal
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = x.to(device)

# Operations
y = x @ W.T                      # matrix multiply
y = x.view(-1, 784)              # reshape
y = x.unsqueeze(0)               # add batch dim
y = torch.cat([a, b], dim=0)     # concatenate

# Autograd
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print(x.grad)  # dy/dx = 2x + 3 = 7.0

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.dropout(self.relu(self.fc2(x)))
        return self.fc3(x)

model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)

        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()

    # Validation
    model.eval()
    with torch.no_grad():
        correct = sum(
            (model(x.to(device)).argmax(1) == y.to(device)).sum().item()
            for x, y in val_loader
        )

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.long)

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

train_loader = DataLoader(
    MyDataset(X_train, y_train),
    batch_size=32, shuffle=True, num_workers=4
)

import tensorflow as tf
from tensorflow import keras
from keras import layers

# Sequential model
model = keras.Sequential([
    layers.Flatten(input_shape=(28, 28)),
    layers.Dense(256, activation="relu"),
    layers.Dropout(0.2),
    layers.Dense(128, activation="relu"),
    layers.Dense(10, activation="softmax"),
])

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"],
)

model.fit(X_train, y_train, epochs=10, batch_size=32,
         validation_split=0.2,
         callbacks=[keras.callbacks.EarlyStopping(patience=3)])

model.evaluate(X_test, y_test)
model.save("my_model.keras")

# PyTorch CNN
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):        # x: (B, 1, 28, 28)
        x = self.pool(torch.relu(self.conv1(x)))  # → (B, 32, 14, 14)
        x = self.pool(torch.relu(self.conv2(x)))  # → (B, 64, 7, 7)
        x = x.view(x.size(0), -1)                  # flatten
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

from torchvision import models

# Load pretrained ResNet, replace final layer
model = models.resnet50(weights="IMAGENET1K_V2")
model.fc = nn.Linear(model.fc.in_features, num_classes)

# Freeze early layers (optional)
for param in model.parameters():
    param.requires_grad = False
for param in model.fc.parameters():
    param.requires_grad = True

# PyTorch LSTM
class LSTMModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_dim * 2, num_classes)

    def forward(self, x):
        x = self.embedding(x)          # (B, seq_len, embed_dim)
        out, (h, c) = self.lstm(x)     # out: (B, seq_len, hidden*2)
        return self.fc(out[:, -1, :])  # last timestep

from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

# Quick inference with pipeline
classifier = pipeline("sentiment-analysis")
result = classifier("This movie was amazing!")
# [{'label': 'POSITIVE', 'score': 0.9998}]

# Text generation
gen = pipeline("text-generation", model="gpt2")
gen("The future of AI is", max_length=50)

# Custom model loading
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

inputs = tokenizer("Hello world", return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    tokenizer=tokenizer,
)
trainer.train()

import mlflow

mlflow.set_experiment("my_experiment")

with mlflow.start_run():
    mlflow.log_param("lr", 1e-3)
    mlflow.log_param("epochs", 10)

    # ... training ...

    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_metric("f1", 0.93)
    mlflow.pytorch.log_model(model, "model")

from fastapi import FastAPI
import torch

app = FastAPI()
model = torch.load("model.pt")
model.eval()

@app.post("/predict")
async def predict(data: dict):
    tensor = torch.tensor(data["features"])
    with torch.no_grad():
        pred = model(tensor)
    return {"prediction": pred.tolist()}

# Check GPU
torch.cuda.is_available()
torch.cuda.device_count()
torch.cuda.get_device_name(0)

# Mixed precision training (2x speed, less memory)
scaler = torch.cuda.amp.GradScaler()
for batch in loader:
    with torch.cuda.amp.autocast():
        output = model(batch)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# Multi-GPU (DataParallel — simple)
model = nn.DataParallel(model)

# Memory tips
torch.cuda.empty_cache()
# Use gradient accumulation for large effective batch size
# Use gradient checkpointing to trade compute for memory