v0.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 MLP(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim, n_layers):
        super().__init__()
        dims = [in_dim] + [hidden_dim] * (n_layers - 1)
        layers = []
        for dim_in, dim_out in zip(dims[:-1], dims[1:]):
            layers.append(nn.Linear(dim_in, dim_out))
            layers.append(nn.ReLU())

        layers.append(nn.Linear(dims[-1], out_dim))  # no ReLU after final logits layer, softmax is backed in "nn.CrossEntropyLoss"

        self.layers = nn.Sequential(*layers)


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

model = MLP(D, H, C, L).to(device)
optimizer = torch.optim.SGD(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}")