What/Why/How of Neural Networks

example.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from tqdm import tqdm


class MLP(nn.Module):
    def __init__(self, feat_in: int=1, feat_out: int=1):
        super(MLP, self).__init__()
        self.feat_in = feat_in
        self.feat_out = feat_out

        self.layer1 = nn.Linear(feat_in, 256)
        self.layer2 = nn.Linear(256, 256)
        self.layer3 = nn.Linear(256, feat_out)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = self.layer3(x)

        return x

class Samples(torch.utils.data.Dataset):
    def __init__(self, sample_pts, labels):
        self.sample_pts = sample_pts
        self.labels = labels

    def __len__(self):
        return self.sample_pts.shape[0]

    def __getitem__(self, idx):
        return self.sample_pts[idx], self.labels[idx]


def main():
    # make data
    x_min, x_max = -10, 10
    sample_pts = torch.Tensor(10000,1).uniform_(x_min, x_max)
    labels = torch.exp(torch.cos(sample_pts)) - 0.125*sample_pts + torch.sin(sample_pts)
    
    # plot actual
    plt.xlim(x_min, x_max)
    plt.ylim(-2, 4)
    plt.scatter(sample_pts, labels)
    plt.savefig(f'plots/actual.png')

    # initialize training
    neural_network = MLP(1, 1)
    print(neural_network)
    optimizer = torch.optim.SGD(neural_network.parameters(), lr=0.001)
    criterion = torch.nn.L1Loss()
    
    samples = Samples(sample_pts, labels)
    samples_loader = torch.utils.data.DataLoader(samples, batch_size=64)
    
    # training
    for epoch in tqdm(range(1, 501), desc=f'Epoch'):
        for sample, label in samples_loader:
            optimizer.zero_grad()

            # predict and update
            prediction = neural_network(sample)
            loss = criterion(prediction, label)
            loss.backward()
            optimizer.step()
            
        # plot prediction
        plt.clf()
        plt.xlim(x_min, x_max)
        plt.ylim(-2, 4)
        plt.title(f'Epoch {epoch}')
        plt.scatter(sample_pts, neural_network(sample_pts).detach())
        plt.savefig(f'plots/epoch_{epoch}.png')


if __name__ == '__main__':
    main()