Neural Networks

import torch
import torch.nn as nn
import torch.nn.functional as F

# As modules
relu  = nn.ReLU()
gelu  = nn.GELU()
# As functions
x_relu = F.relu(x)
x_gelu = F.gelu(x)

linear = nn.Linear(in_features=784, out_features=256, bias=True)
# Parameter count: 784*256 + 256 = 200,960

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim, dropout=0.3):
        super().__init__()
        dims = [input_dim] + hidden_dims + [output_dim]
        layers = []
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i+1]))
            if i < len(dims) - 2:   # no BN/activation on last layer
                layers.append(nn.BatchNorm1d(dims[i+1]))
                layers.append(nn.GELU())
                layers.append(nn.Dropout(dropout))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)

model = MLP(input_dim=128, hidden_dims=[256, 256], output_dim=10)

# Classification
loss_fn = nn.CrossEntropyLoss()   # expects logits (before softmax)
loss = loss_fn(logits, y_true)    # y_true: class indices

# Regression
loss_fn = nn.MSELoss()
loss_fn = nn.SmoothL1Loss(beta=1.0)  # Huber loss

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.block = nn.Sequential(
            nn.Linear(dim, dim), nn.GELU(),
            nn.Linear(dim, dim)
        )
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        return self.norm(x + self.block(x))   # pre-norm or post-norm variant

# Label smoothing
loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

# Mixup (simplified)
lam = np.random.beta(0.4, 0.4)
mixed_x = lam * x1 + (1 - lam) * x2
mixed_y = lam * y1 + (1 - lam) * y2

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# Data
X, y = torch.randn(1000, 64), torch.randint(0, 10, (1000,))
dataset = TensorDataset(X, y)
loader  = DataLoader(dataset, batch_size=64, shuffle=True)

# Model
model = MLP(64, [128, 128], 10)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

# Training loop
for epoch in range(20):
    model.train()
    total_loss = 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        logits = model(xb)
        loss = criterion(logits, yb)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()
    print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}")

# Inference
model.eval()
with torch.no_grad():
    preds = model(X_test.to(device)).argmax(dim=1)

def count_params(model):
    total   = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Total: {total:,}  |  Trainable: {trainable:,}")

count_params(model)