Linear Regression Neural Network with nn.Linear() in PyTorch
Tags: pytorch, linearregression, neuralnetwork, deeplearning
Published: 2024-06-06 01:20:11

```python
import torch
from torch import nn
import matplotlib.pyplot as plt

# Setup device
device = "cuda" if torch.cuda.is_available() else "cpu"
# print(device)

# Create data
weight = 0.7
bias = 0.3

X = torch.arange(start=0, end=1, step=0.02, device=device).unsqueeze(dim=1)
y = weight * X + bias
# print(X[:10], len(X))
# print(y[:10], len(y))

l = int(0.8 * len(X))
X_train, y_train, X_test, y_test = X[:l], y[:l], X[l:], y[l:]
# print(len(X_train), len(y_train), len(X_test), len(y_test))

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_layer = nn.Linear(in_features=1, out_features=1)
    
    def forward(self, x):
        return self.linear_layer(x)

torch.manual_seed(42)

my_model = MyModel().to(device)
# print(my_model, my_model.state_dict())
# print(next(my_model.parameters()).device)
# print(next(my_model.parameters()))

loss_fn = nn.L1Loss()
# loss_fn = nn.MSELoss()

optimizer = torch.optim.SGD(params=my_model.parameters(), lr=0.01)
# optimizer = torch.optim.Adam(params=my_model.parameters(), lr=0.01)

epochs = 100 # Try 0, 50, 100, 150

epoch_count = []
loss_values = []
test_loss_values = []

for epoch in range(epochs):
    my_model.train()
    
    # 1. Calculate predictions
    y_pred = my_model(X_train)
    
    # 2. Calculate loss
    loss = loss_fn(y_pred, y_train)
    
    # 3. Zero out gradient
    optimizer.zero_grad()
    
    # 4. Do backpropagation
    loss.backward()
    
    # 5. Optimize model
    optimizer.step()
    
    # Test
    my_model.eval()
    with torch.inference_mode():
        test_pred = my_model(X_test)
        test_loss = loss_fn(test_pred, y_test)
    if epoch % 10 == 0:
        epoch_count.append(epoch)
        loss_values.append(loss)
        test_loss_values.append(test_loss)
        # print(f"Epoch: {epoch} | Loss: {loss} | Test loss: {test_loss}")

# Visualize
with torch.inference_mode():
    y_pred = my_model(X_test)

def plot_predictions(X_train, y_train, X_test, y_test, predictions=None):
    plt.figure(figsize=[6, 4])
    plt.scatter(X_train, y_train, c='g', s=1, label='Train data')
    plt.scatter(X_test, y_test, c='b', s=3, label='Test data')
    if predictions is not None:
        plt.scatter(X_test, predictions, c='r', s=5, label='Predictions')
    plt.title("Train and test data and predictions")
    plt.legend(prop={'size': 14})

plot_predictions(X_train=X_train.cpu(),
                 y_train=y_train.cpu(),
                 X_test=X_test.cpu(),
                 y_test=y_test.cpu(),
                 predictions=y_pred.cpu())

def plot_loss_curves(epoch_count, loss_values, test_loss_values):
    plt.figure(figsize=[6, 4])
    plt.plot(epoch_count, loss_values, label="Train loss")
    plt.plot(epoch_count, test_loss_values, label="Test loss")
    plt.title("Train and test loss curves")
    plt.ylabel("Loss")
    plt.xlabel("Epochs")
    plt.legend(prop={'size': 14})
    
plot_loss_curves(epoch_count=epoch_count,
                 loss_values=torch.tensor(loss_values).cpu(),
                 test_loss_values=torch.tensor(test_loss_values).cpu())
```

# <`L1Loss()` and `SGD()`>

### `epochs = 0`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/5eqw98uefzyvpu1u35vp.png)

### `epochs = 50`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/sgho3jho85qxouvrre8g.png)

### `epochs = 100`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/qnopfdivr2ddrdj319ep.png)

### `epochs = 150`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/copodpdrid5wmzyovgdz.png)

# <`MSELoss()` and `SGD()`>

### `epochs = 0`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/iv3akejzwwwymnm5ghq6.png)

### `epochs = 50`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/qp8x8zai3bxcqvw3c73d.png)

### `epochs = 100`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9q5fvz7log5hu7s2t1de.png)

### `epochs = 150`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/1ul6yw3unyd0g4dfonc4.png)

# <`L1Loss()` and `Adam()`>

### `epochs = 0`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/e3u8poj5tyjt9ifsvd6f.png)

### `epochs = 50`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/t2uke20ybcfh3ytu41nv.png)

### `epochs = 100`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/2s05lvt3g1l4rik6fhir.png)

### `epochs = 150`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/spkc9f8yv6n6xbuj2zhp.png)

# <`MSELoss()` and `Adam()`>

### `epochs = 0`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/hu7cce7fij6tfgamezwd.png)

### `epochs = 50`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/9pdkhs2vnk7bnltujn17.png)

### `epochs = 100`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/6nadveoioy7xdp4indbp.png)

### `epochs = 150`:

![Image description](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/blj7bp35fyrxpx3doedu.png)