v1.py

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

# ---- config ----
N, D, H, C, L = 4096, 512, 1024, 10, 12
batch_size = 64
epochs = 5
lr = 1e-2
device = "cuda" if torch.cuda.is_available() else "cpu"

# ---- data ----
class DemoDataset(Dataset):
    def __init__(self, n, d, c):
        """
        this method is mostly for setting the path to data folder, split IDs, augemntations, ...

        reading data (images, text files, ...) can happen here if data is small enough to comfortably fit in RAM (an when using few workers)
        each worker will create an instance of this class, meaning it will run this method and create / allocate memory for whatever its implementation do
        """
        self.x = torch.randn(n, d)
        self.y = torch.randint(0, c, (n,))

    # this is used to determin how many batches / samples in an epoch
    def __len__(self):
        return self.x.size(0)

    # this is what's called by a single worker to callect samples for the batch
    def __getitem__(self, idx):
        """
        it's most common and recommended to read / preprocess a data sample in here on the fly
        """
        return self.x[idx], self.y[idx]


dataset = DemoDataset(N, D, C)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# ---- model ----
class MLPBlock(nn.Module):   #& pre-norm residual block
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fc = nn.Linear(dim, dim)

    def forward(self, x):
        return x + torch.relu(self.fc(self.norm(x)))   #& residual add around norm+linear+relu

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, n_layers):
        super().__init__()
        self.in_proj = nn.Linear(in_dim, hidden_dim)

        self.blocks = nn.Sequential(
            *[MLPBlock(hidden_dim) for _ in range(n_layers - 1)]   #& stack of residual blocks
        )
        # we can equtally use nn.ModuleList for self.blocks


        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.in_proj(x) 
        x = self.blocks(x)

        # if we used nn.ModuleList for self.blocks, it's forward would be: 
        # for block in self.blocks:
        #     x = block(x)

        return self.out(x)

model = MLP(D, H, C, L).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr) #&
loss_fn = nn.CrossEntropyLoss()

# ---- train ----
model.train()
for epoch in range(epochs):
    total_loss, seen = 0.0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)

        # by default .backward() add the computed grad to the existing one, so we zero out the old grad each step
        optimizer.zero_grad()
        logits = model(xb)
        loss = loss_fn(logits, yb)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * xb.size(0)
        seen += xb.size(0)

    print(f"epoch {epoch}: loss {total_loss / seen:.4f}")