Pytorch Springboard

pytorch_testing_springboard.py

"""
 A self-contained script in PyTorch that generates a simple dataset from a
 linear model and fits a neural network.

 Can be used as a starting script for trying out different PyTorch functionalities
 """
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


class Net(nn.Module):
    def __init__(self, pp):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(pp, 16)
        self.fc2 = nn.Linear(16, 1)

    def forward(self, t):
        t = F.relu(self.fc1(t))
        t = self.fc2(t)
        return t


def create_data(NN, pp):
    x = torch.randn(NN, pp)
    beta = torch.randn(pp, 1)
    y = torch.mm(x, beta)

    return x, y


def train():
    x, y = create_data(500, 5)

    net = Net(x.shape[1])
    optimizer = optim.Adam(net.parameters(), lr=0.001)
    crit = nn.MSELoss()

    for epoch in range(1000):
        optimizer.zero_grad()

        out = net(x)
        loss = crit(out.squeeze(), y.squeeze())
        loss.backward()

        optimizer.step()

        if epoch % 250 == 0:
            print(loss.item())


if __name__ == "__main__":
    train()