# First, let's import PyTorch
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"GPU available: {torch.cuda.is_available()}")PyTorch version: 2.5.1
GPU available: FalseNN with PyTorch
PyTorch is a framework that handles the heavy lifting of neural networks:
| Component | What PyTorch Does |
|---|---|
| Weights & Biases | Stores them as “tensors” (fancy arrays) that can run on GPUs |
| Forward Pass | Computes outputs from inputs |
| Backward Pass | Automatically calculates gradients for learning |
| Optimisers | Updates weights to improve the model |
Building a Simple Neural Network
In PyTorch, we define a neural network as a Python class. Here’s a simple one with:
- An input layer (takes our data)
nn.Linear(4, 8)- A hidden layer (where the features and rules get chosen)
nn.ReLU()
nn.Linear(8, 12)
nn.ReLU()
nn.Linear(12, 8)
nn.ReLU()- An output layer (gives us predictions)
nn.Linear(8, 3)from torch import nn
class SimpleClassifier(nn.Module):
"""
A simple multi-class classifier
"""
def __init__(self, input_size, hidden_sizes, num_classes):
super().__init__()
# The layers with weights and biases
self.layer_sizes = [input_size, *hidden_sizes, num_classes]
self.layers = nn.ModuleList(
[nn.Linear(i, o) for (i, o) in zip(self.layer_sizes[:-1], self.layer_sizes[1:])]
)
# Activation function (adds non-linearity)
self.relu = nn.ReLU()
def forward(self, x):
"""This defines how data flows through the network"""
# We have layer -> activation -> layer -> activation -> ...
# up until the last layer of the network, for which there is
# no activation
for layer in self.layers[:-1]:
x = layer(x)
x = self.relu(x)
return self.layers[-1](x)
# Create an instance: 4 inputs, 8 hidden neurons, 3 output classes
model = SimpleClassifier(
input_size=4,
hidden_sizes=[8, 12, 8],
num_classes=3,
)Code
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
def plot_model_architecture(model):
layer_sizes = model.layer_sizes
n_layers = len(layer_sizes)
max_neurons = max(layer_sizes)
colors = ["#4C9BE8", "#4CAF50", "#E57373", "#9575CD", "#FFB300"]
layer_names = ["Input"] + [f"Hidden {i}" for i in range(1, n_layers - 1)] + ["Output"]
# Fix figure size to keep aspect ratio balanced
fig_width = n_layers * 2
fig_height = max_neurons * 0.8
fig, ax = plt.subplots(figsize=(fig_width, fig_height))
# Use equal numeric spacing
x_spacing = 3
y_spacing = 1
ax.set_xlim(-1, (n_layers - 1) * x_spacing + 1)
ax.set_ylim(-1, (max_neurons - 1) * y_spacing + 1)
ax.set_aspect("equal") # <-- key fix for round circles
ax.axis("off")
neuron_positions = []
for layer_idx, n_neurons in enumerate(layer_sizes):
color = colors[layer_idx % len(colors)]
x = layer_idx * x_spacing
# Center neurons vertically
total_height = (n_neurons - 1) * y_spacing
y_start = (max_neurons - 1) / 2 - total_height / 2
ys = [y_start + i * y_spacing for i in range(n_neurons)]
neuron_positions.append((x, ys))
for y in ys:
circle = plt.Circle((x, y), 0.3, color=color, zorder=3,
ec="white", linewidth=2)
ax.add_patch(circle)
# Draw connections — slightly thicker edges
for i in range(len(neuron_positions) - 1):
x1, ys1 = neuron_positions[i]
x2, ys2 = neuron_positions[i + 1]
for y1 in ys1:
for y2 in ys2:
ax.plot([x1, x2], [y1, y2], color="gray", alpha=0.3, lw=1.2, zorder=1)
# Labels below each layer
for layer_idx, name in enumerate(layer_names):
x = layer_idx * x_spacing
ax.text(x, -0.8, name, ha="center", va="top", fontsize=11, fontweight="bold")
# Legend
legend_labels = [f"{name} ({n} neurons)" for name, n in zip(layer_names, layer_sizes)]
patches = [mpatches.Patch(color=colors[i % len(colors)], label=legend_labels[i])
for i in range(n_layers)]
ax.legend(handles=patches, loc="upper right", fontsize=9)
plt.title("Neural Network Architecture", fontsize=14, fontweight="bold", pad=20)
plt.tight_layout()
plt.show()
plot_model_architecture(model)
Peeking at the Weights and Biases
The network has randomly initialised weights and biases. These are the numbers that will be adjusted during training.
# Let's look at the parameters in layer 1
print("Layer 1 weights shape:", model.layers[0].weight.shape)
print("Layer 1 biases shape:", model.layers[0].bias.shape)
print("\nActual weight values (randomly initialised):")
print(model.layers[0].weight.data)
# Count total parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"\nTotal trainable parameters: {total_params}")Layer 1 weights shape: torch.Size([8, 4])
Layer 1 biases shape: torch.Size([8])
Actual weight values (randomly initialised):
tensor([[ 0.4132, 0.2714, -0.4756, 0.0829],
[-0.0540, -0.2708, 0.0958, -0.4025],
[ 0.2676, 0.0328, -0.4363, -0.4139],
[-0.1430, -0.4163, 0.2974, -0.2442],
[-0.2699, -0.2285, -0.0606, 0.1280],
[ 0.3851, -0.3586, 0.0128, 0.1410],
[ 0.4227, 0.4052, 0.3390, -0.2431],
[-0.0537, 0.4721, -0.0776, -0.4492]])
Total trainable parameters: 279Making a Prediction (Forward Pass)
Let’s pass some fake data through our network. The output will be meaningless (we haven’t trained it yet!), but this shows how data flows through.
# Create some fake input data: 1 sample with 4 features
fake_input = torch.tensor([[5.1, 3.5, 1.4, 0.2]])
# Pass it through the network
output = model(fake_input)
print("Input shape:", fake_input.shape)
print("Output shape:", output.shape)
print("\nRaw output (logits):", output)
# Convert to probabilities using softmax
probabilities = torch.softmax(output, dim=1)
print("\nAs probabilities:", probabilities)
print(f"\nPredicted class: {probabilities.argmax().item()}")Input shape: torch.Size([1, 4])
Output shape: torch.Size([1, 3])
Raw output (logits): tensor([[-0.0699, -0.0752, -0.3735]], grad_fn=<AddmmBackward0>)
As probabilities: tensor([[0.3659, 0.3640, 0.2701]], grad_fn=<SoftmaxBackward0>)
Predicted class: 0The prediction above was random because our network hasn’t learned anything yet. Training is where the magic happens!