Utilize PyTorch in Linear Regression.

linear_regression.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torch import nn


class RandomLinearDataset(Dataset):
    def __init__(self, x_min: float, x_max: float, n_samples: int, scale: float = 100):
        k, b = [round(x.item() * scale, 4) for x in torch.rand(2)]
        x = torch.unsqueeze(torch.linspace(x_min, x_max, n_samples), dim=1)
        noise = torch.randn(x.size())
        y = k * x + b + noise

        self.n_samples = n_samples
        self.x, self.y = x, y
        self.k, self.b = k, b

    def __len__(self):
        return self.n_samples

    def __getitem__(self, index) -> tuple:
        return self.y[index], self.x[index]


class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.nn_stack = nn.Linear(1, 1)

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


class Couch:
    device = "cuda" if torch.cuda.is_available() else "cpu"

    def __init__(self, model: nn.Module):
        self.model = model.to(self.device)
        self.loss_fn = nn.MSELoss()
        self.optimizer = torch.optim.SGD(model.parameters(), lr=.01)

    def train(self, data_loader: DataLoader, n_epochs: int):
        self.model.train()
        for t in range(n_epochs):
            # print(" Epoch {0:03} ".format(t + 1).join(["=" * 16] * 2))
            for batch_index, (ys, xs) in enumerate(data_loader):
                predict = self.model(xs)
                if predict.isinf().any():
                    raise RuntimeError("Predict inf!")
                if predict.isnan().any():
                    raise RuntimeError("Predict nan!")
                loss = self.loss_fn(predict, ys)

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
        return self

    def test(self, data_loader: DataLoader):
        self.model.eval()
        print(" Test ".join(["=" * 16] * 2))

        ds = data_loader.dataset
        ys, xs = ds.y, ds.x

        predict = self.model(xs)
        # loss = self.loss_fn(predict, ys)

        print("Dataset parameters:",
              "\n\tk=", ds.k,
              "\n\tb=", ds.b)
        print("Model parameters:",
              *["\n\t{}= {}".format(k, v.item())
                for k, v in self.model.named_parameters()])

        plt.plot(xs.numpy(), predict.detach().numpy(), c='r')
        plt.scatter(xs.numpy(), ys.numpy())
        plt.show()

    def save(self, path="model.pth"):
        torch.save(self.model.state_dict(), path)
        print(f"PyTorch Model State saved to {path}.")

    def train_and_test(self, train_data_loader, n_epochs, test_data_loader):
        self.train(train_data_loader, n_epochs=n_epochs)
        self.test(test_data_loader)


if __name__ == "__main__":
    couch = Couch(LinearRegression())
    dataset = RandomLinearDataset(n_samples=100, x_max=1, x_min=-1, scale=10)
    couch.train_and_test(
        DataLoader(dataset, shuffle=True), 20,
        DataLoader(dataset)
    )