Simple PyTorch Boilerplate

pytorch_boilerplate.py

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

class SampleDataset(Dataset):
    def __init__(self, size, coef_1, coef_2):
        self.size = size
        self.coef_1 = coef_1
        self.coef_2 = coef_2
        
    def __len__(self):
        return self.size
    
    def __getitem__(self, idx):
        n1 = random.randint(1, 1000)
        n2 = random.randint(1, 1000)
        feature = torch.Tensor([n1, n2])
        output = torch.Tensor([n1*self.coef_1 + n2*self.coef_2])
        return feature, output
        
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(2, 1)
        
    def forward(self, x):
        out = self.fc1(x)
        return out

dataset = SampleDataset(size=10000, coef_1=3, coef_2=4)
dataloader = DataLoader(dataset, batch_size=32)

model = Net()
optimizer = Adam(model.parameters())
criterion = nn.MSELoss()

EPOCHS = 30
for epoch in range(EPOCHS):
    epoch_loss = 0
    for i, (features, output) in enumerate(dataloader):
        optimizer.zero_grad()
        prediction = model(features)
        loss = criterion(prediction, output)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    normalized_loss = epoch_loss/len(output)
    print(f'EPOCH - {epoch + 1} | Loss - {normalized_loss}')