# Install the mnist1d package from GitHub
!pip install git+https://github.com/greydanus/mnist1d

Collecting git+https://github.com/greydanus/mnist1d
  Cloning https://github.com/greydanus/mnist1d to c:\users\user\appdata\local\temp\pip-req-build-qb2u6in6
  Resolved https://github.com/greydanus/mnist1d to commit 7878d96082abd200c546a07a4101fa90b30fdf7e
  Installing build dependencies: started
  Installing build dependencies: finished with status 'done'
  Getting requirements to build wheel: started
  Getting requirements to build wheel: finished with status 'done'
  Preparing metadata (pyproject.toml): started
  Preparing metadata (pyproject.toml): finished with status 'done'
Requirement already satisfied: requests in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from mnist1d==0.0.2.post16) (2.32.5)
Requirement already satisfied: numpy in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from mnist1d==0.0.2.post16) (2.3.3)
Requirement already satisfied: matplotlib in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from mnist1d==0.0.2.post16) (3.10.7)
Requirement already satisfied: scipy in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from mnist1d==0.0.2.post16) (1.17.0)
Requirement already satisfied: contourpy>=1.0.1 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (1.3.3)
Requirement already satisfied: cycler>=0.10 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (4.60.1)
Requirement already satisfied: kiwisolver>=1.3.1 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (1.4.9)
Requirement already satisfied: packaging>=20.0 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (25.0)
Requirement already satisfied: pillow>=8 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (11.3.0)
Requirement already satisfied: pyparsing>=3 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (3.2.5)
Requirement already satisfied: python-dateutil>=2.7 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from matplotlib->mnist1d==0.0.2.post16) (2.9.0.post0)
Requirement already satisfied: charset_normalizer<4,>=2 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from requests->mnist1d==0.0.2.post16) (3.4.4)
Requirement already satisfied: idna<4,>=2.5 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from requests->mnist1d==0.0.2.post16) (3.11)
Requirement already satisfied: urllib3<3,>=1.21.1 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from requests->mnist1d==0.0.2.post16) (2.5.0)
Requirement already satisfied: certifi>=2017.4.17 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from requests->mnist1d==0.0.2.post16) (2025.10.5)
Requirement already satisfied: six>=1.5 in c:\users\user\documents\github\simpe-ai\venv\lib\site-packages (from python-dateutil>=2.7->matplotlib->mnist1d==0.0.2.post16) (1.17.0)
Building wheels for collected packages: mnist1d
  Building wheel for mnist1d (pyproject.toml): started
  Building wheel for mnist1d (pyproject.toml): finished with status 'done'
  Created wheel for mnist1d: filename=mnist1d-0.0.2.post16-py3-none-any.whl size=14757 sha256=28248f10e8e57530856d80f51f156c101fb57a062e664d4f84506608d5825f11
  Stored in directory: C:\Users\user\AppData\Local\Temp\pip-ephem-wheel-cache-f3qsrxdj\wheels\18\40\b6\29381fee9b4c80fdbc304d52bb065a7286bbcca5ca2b8737c0
Successfully built mnist1d
Installing collected packages: mnist1d
Successfully installed mnist1d-0.0.2.post16

  Running command git clone --filter=blob:none --quiet https://github.com/greydanus/mnist1d 'C:\Users\user\AppData\Local\Temp\pip-req-build-qb2u6in6'

[notice] A new release of pip is available: 25.0.1 -> 26.0.1
[notice] To update, run: python.exe -m pip install --upgrade pip

import numpy as np
import os
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
from torch.optim.lr_scheduler import StepLR
import matplotlib.pyplot as plt
import mnist1d
import random

args = mnist1d.data.get_dataset_args()
data = mnist1d.data.get_dataset(
    args, path='./mnist1d_data.pkl', download=False, regenerate=False)


# print data size and shape
print(f"Training data size: {data['x'].shape}")
print(f"Training labels size: {data['y'].shape}")
print(f"Test data size: {data['x_test'].shape}")
print(f"Test labels size: {data['y_test'].shape}")

Successfully loaded data from ./mnist1d_data.pkl
Training data size: (4000, 40)
Training labels size: (4000,)
Test data size: (1000, 40)
Test labels size: (1000,)

# Load in the data
train_data_x = data['x'].transpose()
train_data_y = data['y']
val_data_x = data['x_test'].transpose()
val_data_y = data['y_test']

Note: We are testing on the MNIST-1D, which contains a signal pattern (pattern of ocillations) for each digit rather than a 28x28 pixle grid like the orignal MNIST-2D. In this case, each number can be represented linearly like the below example:

# display a random sample from the training data
idx = random.randint(0, data['x'].shape[0] - 1)
plt.plot(data['x'][idx])
plt.title(f"Label: {data['y'][idx]}")
plt.show()

The network then should mapp these linear oscillations (x_input) to a digit number (y_output).

From the shape of the data, we saw that there are 40 input dimensions and 10 output dimensions (Each diget represented by 10 values).

class ResidualNetwork(torch.nn.Module):
    def __init__(self, input_size, output_size, hidden_size=100):
      super(ResidualNetwork, self).__init__()
      self.linear1 = nn.Linear(input_size, hidden_size)
      self.linear2 = nn.Linear(hidden_size, hidden_size)
      self.linear3 = nn.Linear(hidden_size, hidden_size)
      self.linear4 = nn.Linear(hidden_size, output_size)
      print("Initialized ResidualNetwork model with {} parameters".format(self.count_params()))
    
    def count_params(self):
        return sum([p.view(-1).shape[0] for p in self.parameters()])

    def forward(self, x):
        h1 = self.linear1(x)
        h2 = h1 + self.linear2(h1.relu())
        h3 = h2 + self.linear3(h2.relu())
        return self.linear4(h3)

# He initialization of weights
def weights_init(layer_in):
  if isinstance(layer_in, nn.Linear):
    nn.init.kaiming_uniform_(layer_in.weight)
    layer_in.bias.data.fill_(0.0)

# define the model, loss function, and optimizer
model = ResidualNetwork(input_size=40, output_size=10, hidden_size=100)
# apply the weights initialization
model.apply(weights_init)

# define the loss function and optimizer
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# object that decreases learning rate by half every 20 epochs
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)

Initialized ResidualNetwork model with 25310 parameters

# convert data to torch tensors
x_train = torch.tensor(train_data_x.transpose().astype('float32'))
y_train = torch.tensor(train_data_y).long()
x_val = torch.tensor(val_data_x.transpose().astype('float32'))
y_val = torch.tensor(val_data_y).long()

# Load the data into PyTorch DataLoaders
data_loader = DataLoader(TensorDataset(
    x_train, y_train), batch_size=100, shuffle=True, worker_init_fn=np.random.seed(1))

# loop over the dataset n_epoch times
n_epoch = 100
# store the loss and the % correct at each epoch
losses_train = np.zeros((n_epoch))
errors_train = np.zeros((n_epoch))
losses_val = np.zeros((n_epoch))
errors_val = np.zeros((n_epoch))


for epoch in range(n_epoch):
    
    # ----- TRAIN LOOP -----
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    # loop over the batches of the dataset
    for i, (x_batch, y_batch) in enumerate(data_loader):
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward pass -- calculate model output
        pred = model(x_batch)
        # compute the loss
        loss = loss_function(pred, y_batch)
        # backward pass
        loss.backward()
        # SGD update
        optimizer.step()
        
        # update running loss and accuracy
        running_loss += loss.item()
        _, predicted = torch.max(pred.data, 1)
        total += y_batch.size(0)
        correct += (predicted == y_batch).sum().item()
    
    # store the average loss and % correct for this epoch
    losses_train[epoch] = running_loss / len(data_loader)
    errors_train[epoch] = 1 - correct / total
       
    # step the learning rate scheduler to decrease the learning rate if necessary
    scheduler.step()
    
    # ----- VALIDATION LOOP -----
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        pred = model(x_val)
        loss = loss_function(pred, y_val)
        running_loss += loss.item()
        _, predicted = torch.max(pred.data, 1)
        total += y_val.size(0)
        correct += (predicted == y_val).sum().item()
    
    # store the average loss and % correct for this epoch
    losses_val[epoch] = running_loss / len(data_loader)
    errors_val[epoch] = 1 - correct / total
    
    
    print(f"Epoch {epoch+1}/{n_epoch} - Train Loss: {losses_train[epoch]:.4f}, Train Error: {errors_train[epoch]:.4f}, Val Loss: {losses_val[epoch]:.4f}, Val Error: {errors_val[epoch]:.4f}")

Epoch 1/100 - Train Loss: 3.9910, Train Error: 0.8438, Val Loss: 0.0675, Val Error: 0.7970
Epoch 2/100 - Train Loss: 2.2723, Train Error: 0.7548, Val Loss: 0.0563, Val Error: 0.7570
Epoch 3/100 - Train Loss: 2.0017, Train Error: 0.7218, Val Loss: 0.0522, Val Error: 0.7350
Epoch 4/100 - Train Loss: 1.8622, Train Error: 0.6947, Val Loss: 0.0499, Val Error: 0.7270
Epoch 5/100 - Train Loss: 1.7781, Train Error: 0.6805, Val Loss: 0.0478, Val Error: 0.7150
Epoch 6/100 - Train Loss: 1.7095, Train Error: 0.6518, Val Loss: 0.0463, Val Error: 0.7030
Epoch 7/100 - Train Loss: 1.6586, Train Error: 0.6378, Val Loss: 0.0454, Val Error: 0.6990
Epoch 8/100 - Train Loss: 1.6174, Train Error: 0.6155, Val Loss: 0.0444, Val Error: 0.6730
Epoch 9/100 - Train Loss: 1.5796, Train Error: 0.6018, Val Loss: 0.0435, Val Error: 0.6600
Epoch 10/100 - Train Loss: 1.5401, Train Error: 0.5825, Val Loss: 0.0428, Val Error: 0.6550
Epoch 11/100 - Train Loss: 1.5076, Train Error: 0.5733, Val Loss: 0.0425, Val Error: 0.6470
Epoch 12/100 - Train Loss: 1.5010, Train Error: 0.5725, Val Loss: 0.0424, Val Error: 0.6490
Epoch 13/100 - Train Loss: 1.4976, Train Error: 0.5695, Val Loss: 0.0424, Val Error: 0.6450
Epoch 14/100 - Train Loss: 1.4944, Train Error: 0.5717, Val Loss: 0.0423, Val Error: 0.6450
Epoch 15/100 - Train Loss: 1.4920, Train Error: 0.5677, Val Loss: 0.0422, Val Error: 0.6460
Epoch 16/100 - Train Loss: 1.4887, Train Error: 0.5682, Val Loss: 0.0422, Val Error: 0.6420
Epoch 17/100 - Train Loss: 1.4862, Train Error: 0.5685, Val Loss: 0.0421, Val Error: 0.6450
Epoch 18/100 - Train Loss: 1.4834, Train Error: 0.5662, Val Loss: 0.0420, Val Error: 0.6430
Epoch 19/100 - Train Loss: 1.4808, Train Error: 0.5627, Val Loss: 0.0420, Val Error: 0.6450
Epoch 20/100 - Train Loss: 1.4782, Train Error: 0.5620, Val Loss: 0.0419, Val Error: 0.6430
Epoch 21/100 - Train Loss: 1.4745, Train Error: 0.5640, Val Loss: 0.0419, Val Error: 0.6430
Epoch 22/100 - Train Loss: 1.4742, Train Error: 0.5640, Val Loss: 0.0419, Val Error: 0.6440
Epoch 23/100 - Train Loss: 1.4739, Train Error: 0.5633, Val Loss: 0.0419, Val Error: 0.6450
Epoch 24/100 - Train Loss: 1.4736, Train Error: 0.5630, Val Loss: 0.0419, Val Error: 0.6440
Epoch 25/100 - Train Loss: 1.4733, Train Error: 0.5630, Val Loss: 0.0419, Val Error: 0.6430
Epoch 26/100 - Train Loss: 1.4731, Train Error: 0.5625, Val Loss: 0.0419, Val Error: 0.6430
Epoch 27/100 - Train Loss: 1.4728, Train Error: 0.5620, Val Loss: 0.0419, Val Error: 0.6410
Epoch 28/100 - Train Loss: 1.4725, Train Error: 0.5613, Val Loss: 0.0419, Val Error: 0.6410
Epoch 29/100 - Train Loss: 1.4723, Train Error: 0.5620, Val Loss: 0.0419, Val Error: 0.6410
Epoch 30/100 - Train Loss: 1.4720, Train Error: 0.5613, Val Loss: 0.0419, Val Error: 0.6410
Epoch 31/100 - Train Loss: 1.4716, Train Error: 0.5608, Val Loss: 0.0419, Val Error: 0.6410
Epoch 32/100 - Train Loss: 1.4716, Train Error: 0.5610, Val Loss: 0.0419, Val Error: 0.6410
Epoch 33/100 - Train Loss: 1.4715, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 34/100 - Train Loss: 1.4715, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 35/100 - Train Loss: 1.4715, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 36/100 - Train Loss: 1.4715, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 37/100 - Train Loss: 1.4714, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 38/100 - Train Loss: 1.4714, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 39/100 - Train Loss: 1.4714, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 40/100 - Train Loss: 1.4714, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 41/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 42/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 43/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 44/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 45/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 46/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 47/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 48/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 49/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 50/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 51/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 52/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 53/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 54/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 55/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 56/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 57/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 58/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 59/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 60/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 61/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 62/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 63/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 64/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 65/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 66/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 67/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 68/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 69/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 70/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 71/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 72/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 73/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 74/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 75/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 76/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 77/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 78/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 79/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 80/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 81/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 82/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 83/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 84/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 85/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 86/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 87/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 88/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 89/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 90/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 91/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 92/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 93/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 94/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 95/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 96/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 97/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 98/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 99/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410
Epoch 100/100 - Train Loss: 1.4713, Train Error: 0.5605, Val Loss: 0.0419, Val Error: 0.6410

import matplotlib.pyplot as plt

epochs = np.arange(1, n_epoch + 1)

# ----- LOSS PLOT -----
plt.figure()
plt.plot(epochs, losses_train, label="Train Loss")
plt.plot(epochs, losses_val, label="Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training and Validation Loss")
plt.legend()
plt.grid(True)
plt.show()

# ----- ERROR PLOT -----
plt.figure()
plt.plot(epochs, errors_train, label="Train Error")
plt.plot(epochs, errors_val, label="Validation Error")
plt.xlabel("Epoch")
plt.ylabel("Error")
plt.title("Training and Validation Error")
plt.legend()
plt.grid(True)
plt.show()

Training without Residual Connection

class Network(torch.nn.Module):
    def __init__(self, input_size, output_size, hidden_size=100):
      super(Network, self).__init__()
      self.linear1 = nn.Linear(input_size, hidden_size)
      self.linear2 = nn.Linear(hidden_size, hidden_size)
      self.linear3 = nn.Linear(hidden_size, hidden_size)
      self.linear4 = nn.Linear(hidden_size, output_size)
      print("Initialized model with {} parameters".format(self.count_params()))
    
    def count_params(self):
        return sum([p.view(-1).shape[0] for p in self.parameters()])

    def forward(self, x):
        h1 = self.linear1(x)
        h2 = self.linear2(h1.relu())
        h3 = self.linear3(h2.relu())
        return self.linear4(h3)


# define the model, loss function, and optimizer
model = Network(input_size=40, output_size=10, hidden_size=100)
# apply the weights initialization
model.apply(weights_init)

# define the loss function and optimizer
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# object that decreases learning rate by half every 20 epochs
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)


# loop over the dataset n_epoch times
n_epoch = 100
# store the loss and the % correct at each epoch
losses_train = np.zeros((n_epoch))
errors_train = np.zeros((n_epoch))
losses_val = np.zeros((n_epoch))
errors_val = np.zeros((n_epoch))


for epoch in range(n_epoch):

    # ----- TRAIN LOOP -----
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    # loop over the batches of the dataset
    for i, (x_batch, y_batch) in enumerate(data_loader):
        # zero the parameter gradients
        optimizer.zero_grad()
        # forward pass -- calculate model output
        pred = model(x_batch)
        # compute the loss
        loss = loss_function(pred, y_batch)
        # backward pass
        loss.backward()
        # SGD update
        optimizer.step()

        # update running loss and accuracy
        running_loss += loss.item()
        _, predicted = torch.max(pred.data, 1)
        total += y_batch.size(0)
        correct += (predicted == y_batch).sum().item()

    # store the average loss and % correct for this epoch
    losses_train[epoch] = running_loss / len(data_loader)
    errors_train[epoch] = 1 - correct / total

    # step the learning rate scheduler to decrease the learning rate if necessary
    scheduler.step()

    # ----- VALIDATION LOOP -----
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        pred = model(x_val)
        loss = loss_function(pred, y_val)
        running_loss += loss.item()
        _, predicted = torch.max(pred.data, 1)
        total += y_val.size(0)
        correct += (predicted == y_val).sum().item()

    # store the average loss and % correct for this epoch
    losses_val[epoch] = running_loss / len(data_loader)
    errors_val[epoch] = 1 - correct / total

    print(
        f"Epoch {epoch+1}/{n_epoch} - Train Loss: {losses_train[epoch]:.4f}, Train Error: {errors_train[epoch]:.4f}, Val Loss: {losses_val[epoch]:.4f}, Val Error: {errors_val[epoch]:.4f}")
    
    import matplotlib.pyplot as plt

epochs = np.arange(1, n_epoch + 1)

# ----- LOSS PLOT -----
plt.figure()
plt.plot(epochs, losses_train, label="Train Loss")
plt.plot(epochs, losses_val, label="Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training and Validation Loss")
plt.legend()
plt.grid(True)
plt.show()

# ----- ERROR PLOT -----
plt.figure()
plt.plot(epochs, errors_train, label="Train Error")
plt.plot(epochs, errors_val, label="Validation Error")
plt.xlabel("Epoch")
plt.ylabel("Error")
plt.title("Training and Validation Error")
plt.legend()
plt.grid(True)
plt.show()

Initialized model with 25310 parameters
Epoch 1/100 - Train Loss: 2.8043, Train Error: 0.8780, Val Loss: 0.0609, Val Error: 0.8320
Epoch 2/100 - Train Loss: 2.2618, Train Error: 0.7940, Val Loss: 0.0546, Val Error: 0.8020
Epoch 3/100 - Train Loss: 2.0685, Train Error: 0.7448, Val Loss: 0.0512, Val Error: 0.7650
Epoch 4/100 - Train Loss: 1.9519, Train Error: 0.7295, Val Loss: 0.0491, Val Error: 0.7390
Epoch 5/100 - Train Loss: 1.8713, Train Error: 0.7037, Val Loss: 0.0474, Val Error: 0.7300
Epoch 6/100 - Train Loss: 1.8097, Train Error: 0.6912, Val Loss: 0.0463, Val Error: 0.7250
Epoch 7/100 - Train Loss: 1.7641, Train Error: 0.6743, Val Loss: 0.0452, Val Error: 0.7090
Epoch 8/100 - Train Loss: 1.7258, Train Error: 0.6613, Val Loss: 0.0445, Val Error: 0.7060
Epoch 9/100 - Train Loss: 1.6945, Train Error: 0.6412, Val Loss: 0.0437, Val Error: 0.6990
Epoch 10/100 - Train Loss: 1.6652, Train Error: 0.6345, Val Loss: 0.0435, Val Error: 0.6920
Epoch 11/100 - Train Loss: 1.6444, Train Error: 0.6295, Val Loss: 0.0433, Val Error: 0.6930
Epoch 12/100 - Train Loss: 1.6393, Train Error: 0.6295, Val Loss: 0.0432, Val Error: 0.6940
Epoch 13/100 - Train Loss: 1.6361, Train Error: 0.6228, Val Loss: 0.0431, Val Error: 0.6910
Epoch 14/100 - Train Loss: 1.6340, Train Error: 0.6215, Val Loss: 0.0430, Val Error: 0.6890
Epoch 15/100 - Train Loss: 1.6314, Train Error: 0.6202, Val Loss: 0.0430, Val Error: 0.6900
Epoch 16/100 - Train Loss: 1.6289, Train Error: 0.6202, Val Loss: 0.0429, Val Error: 0.6900
Epoch 17/100 - Train Loss: 1.6267, Train Error: 0.6170, Val Loss: 0.0429, Val Error: 0.6900
Epoch 18/100 - Train Loss: 1.6243, Train Error: 0.6173, Val Loss: 0.0428, Val Error: 0.6880
Epoch 19/100 - Train Loss: 1.6222, Train Error: 0.6178, Val Loss: 0.0428, Val Error: 0.6880
Epoch 20/100 - Train Loss: 1.6202, Train Error: 0.6160, Val Loss: 0.0428, Val Error: 0.6860
Epoch 21/100 - Train Loss: 1.6177, Train Error: 0.6158, Val Loss: 0.0428, Val Error: 0.6850
Epoch 22/100 - Train Loss: 1.6175, Train Error: 0.6160, Val Loss: 0.0427, Val Error: 0.6850
Epoch 23/100 - Train Loss: 1.6173, Train Error: 0.6150, Val Loss: 0.0427, Val Error: 0.6850
Epoch 24/100 - Train Loss: 1.6171, Train Error: 0.6148, Val Loss: 0.0427, Val Error: 0.6850
Epoch 25/100 - Train Loss: 1.6168, Train Error: 0.6135, Val Loss: 0.0427, Val Error: 0.6860
Epoch 26/100 - Train Loss: 1.6166, Train Error: 0.6140, Val Loss: 0.0427, Val Error: 0.6850
Epoch 27/100 - Train Loss: 1.6164, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6850
Epoch 28/100 - Train Loss: 1.6162, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 29/100 - Train Loss: 1.6160, Train Error: 0.6135, Val Loss: 0.0427, Val Error: 0.6860
Epoch 30/100 - Train Loss: 1.6158, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 31/100 - Train Loss: 1.6155, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 32/100 - Train Loss: 1.6155, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 33/100 - Train Loss: 1.6155, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 34/100 - Train Loss: 1.6155, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 35/100 - Train Loss: 1.6154, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 36/100 - Train Loss: 1.6154, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 37/100 - Train Loss: 1.6154, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 38/100 - Train Loss: 1.6154, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 39/100 - Train Loss: 1.6153, Train Error: 0.6128, Val Loss: 0.0427, Val Error: 0.6860
Epoch 40/100 - Train Loss: 1.6153, Train Error: 0.6133, Val Loss: 0.0427, Val Error: 0.6860
Epoch 41/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 42/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 43/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 44/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 45/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 46/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 47/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 48/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 49/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 50/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 51/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 52/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 53/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 54/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 55/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 56/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 57/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 58/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 59/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 60/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 61/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 62/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 63/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 64/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 65/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 66/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 67/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 68/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 69/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 70/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 71/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 72/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 73/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 74/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 75/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 76/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 77/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 78/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 79/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 80/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 81/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 82/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 83/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 84/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 85/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 86/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 87/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 88/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 89/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 90/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 91/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 92/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 93/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 94/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 95/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 96/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 97/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 98/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 99/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860
Epoch 100/100 - Train Loss: 1.6153, Train Error: 0.6130, Val Loss: 0.0427, Val Error: 0.6860