PyTorch nn.Linear from First Principles

Hand-verify nn.Linear math, plot activation functions, and trace a full forward pass.

Introduction

30 Second Summary

Neural networks process millions of numbers through layers of math, but most tutorials skip over what actually happens to those numbers at each step. That gap turns debugging into guesswork.

In this project, you will build a hands-on Google Colab notebook where you compute every operation inside PyTorch's nn.Linear and activation functions by hand, proving to yourself there is no magic hidden inside these layers.

What You'll Build

Picture opening your notebook, passing a tensor through a network, and being able to point at any intermediate value and say exactly how it got there.

By the end of this project, you'll have:

  • A hand-verified proof that nn.Linear just computes y = xW^T + b, with your manual matrix math matching PyTorch's output to floating point precision.
  • matplotlib plots of ReLU, Sigmoid, Tanh, and LeakyReLU with their derivatives, showing exactly where each function saturates or dies.
  • A complete forward pass through a 2-layer network where you can trace every intermediate tensor from input to output, verifying each value against your own calculations.
  • Secret Mission: Build your own custom Linear layer from scratch using nn.Parameter and prove it produces identical results to PyTorch's built-in version.

Are there any prerequisites?

You need basic Python scripting experience and a general understanding of matrix multiplication (rows times columns). No prior PyTorch or deep learning experience is required.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters. This is your moment to commit to the goal before any hands-on work begins.

Set Up PyTorch in Colab

Before you explore PyTorch's neural network internals, you need a working environment where every number is reproducible. Google Colab provides a free, browser-based notebook with PyTorch pre-installed.

In this step, you'll confirm PyTorch is available, import the libraries you need, and set a manual seed so every random tensor is identical across runs. Reproducibility is essential because you will hand-verify every number PyTorch produces throughout this project.

In this step, get ready to:
  • Confirm PyTorch is available in Google Colab.
  • Set a manual seed for reproducible results.
  • Verify reproducibility by running the same cell twice.
Open a Google Colab notebook and confirm PyTorch
  • Open Google Colab by navigating to colab.research.google.com in your browser.
  • Sign in with your Google account if prompted.
  • Click New notebook to create a fresh notebook.
  • In the first code cell, confirm PyTorch is pre-installed by running this code:
import torch
print(f"PyTorch version: {torch.__version__}")

You should see a version number printed below the cell, such as 2.10.0.

Seeing an import error?

Google Colab comes with PyTorch pre-installed on every runtime. If you see ModuleNotFoundError, try clicking Runtime in the top menu, then Disconnect and delete runtime, then reconnect by running the cell again.

help me fix a PyTorch import error in Google Colab

Set up imports and a manual seed

A manual seed forces PyTorch's random number generator to start from the same point every time. This means any random tensor you create will have identical values across runs, letting you hand-verify calculations.

  • Replace your cell's contents with the full setup code below and run it:
# Cell 1: Setup
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

torch.manual_seed(42)
print(f"PyTorch version: {torch.__version__}")
print(f"Reproducibility check: {torch.randn(2)}")

What does this code do?

  • torch is the core PyTorch library. torch.nn provides neural network layers, and torch.nn.functional provides activation functions you'll use in later steps.
  • matplotlib.pyplot is the plotting library you'll use to visualize activation functions.
  • torch.manual_seed(42) locks the random number generator to a fixed starting point. Every call to torch.randn() after this seed will produce the same values.

You should see two lines of output: the PyTorch version number and a tensor with two random values.

Not seeing any output?

Make sure you clicked the play button on the left side of the cell (or pressed Shift+Enter) to run it. If you see a connection error, click Connect in the top-right corner of Colab to start a runtime.

help me run a cell in Google Colab

Verify reproducibility by re-running the cell

Before you run the cell again, predict: will the random tensor show the same two values, or different ones?

  • Run the cell a second time by pressing Shift+Enter or clicking the play button.
  • Compare the tensor values to your first run.

The random tensor prints the exact same values both times. This confirms that torch.manual_seed(42) resets the generator every time the cell runs, making every number in this notebook reproducible and hand-verifiable.

Why does the seed reset on every run?

Because torch.manual_seed(42) sits at the top of the cell, it executes before torch.randn(2) every time you re-run. The generator always starts from the same state, producing identical output regardless of how many times you run the cell.

✔️ Awesome, I've got everything!

Great. Double check you've saved your notebook. Colab auto-saves, but you can press Ctrl+S (Windows) or Cmd+S (macOS) to be sure.

ⓧ I'd like to double check the full code

Here is the complete Cell 1 code. Your cell should match this exactly:

# Cell 1: Setup
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

torch.manual_seed(42)
print(f"PyTorch version: {torch.__version__}")
print(f"Reproducibility check: {torch.randn(2)}")

Your environment is locked and loaded. Every random number from here on is reproducible. Next up, you will get hands-on with PyTorch's core data structure — the tensor — and master the operations that power every neural network layer.

Explore Tensors, Dtypes, and Core Operations

Your environment is ready and every random number is reproducible. Before you touch neural network layers, you need fluency with the object that holds all the data flowing through them: the tensor.

A tensor is PyTorch's version of a multi-dimensional array. Every weight matrix, every input batch, and every gradient is a tensor. In this step, you will create tensors, inspect their properties, reshape them, and perform the exact matrix math that nn.Linear uses internally.

In this step, get ready to:
  • Create tensors using different factory functions and inspect their shape, dtype, and device.
  • Reshape and slice tensors to understand how PyTorch organizes data in memory.
  • Perform matrix multiplication and broadcasting that mirrors what nn.Linear does under the hood.
Create tensors and inspect their properties

PyTorch provides several factory functions for creating tensors. Each tensor carries three key properties: its shape (dimensions), dtype (data type), and device (CPU or GPU).

  • Create a new code cell below your setup cell. Explore tensor creation and properties by pasting this code:
# Cell 2: Tensor Basics
# Create tensors using different factory functions
scalar = torch.tensor(3.14)
vector = torch.tensor([1.0, 2.0, 3.0])
matrix = torch.zeros(2, 3)
random_t = torch.randn(2, 3)

print("=== Tensor Creation ===")
print(f"scalar: value={scalar}, shape={scalar.shape}, dtype={scalar.dtype}")
print(f"vector: {vector}, shape={vector.shape}")
print(f"matrix (zeros):\n{matrix}")
print(f"random:\n{random_t}")
print(f"device: {random_t.device}")

# Dtypes and casting
int_tensor = torch.tensor([1, 2, 3])
float_tensor = int_tensor.float()
bool_tensor = torch.tensor([True, False, True])

print("\n=== Dtypes and Casting ===")
print(f"int_tensor dtype: {int_tensor.dtype}")
print(f"After .float(): {float_tensor.dtype}")
print(f"bool_tensor dtype: {bool_tensor.dtype}")
print(f"bool as float: {bool_tensor.float()}")

What does this code do?

torch.tensor() creates a tensor from a Python value or list. torch.zeros() and torch.randn() create tensors of a given shape filled with zeros or random values from a normal distribution.

Every tensor has a dtype. PyTorch defaults to float32 for decimal numbers and int64 for integers. The .float() method casts any tensor to float32. Neural networks require float tensors for their math, so you will use this cast often.

  • Run the cell.
  • You should see each tensor printed with its value and shape. The scalar has shape torch.Size([]) (zero dimensions), the vector has torch.Size([3]), and the matrix has torch.Size([2, 3]). The dtype section shows how .float() converts integers to torch.float32.

Seeing a NameError for torch?

Make sure you ran Cell 1 (your setup cell) first in this session. The import torch statement must execute before any other cell can use torch.

Help me fix a NameError when creating tensors in Colab.

Reshape and slice tensor data

Neural networks constantly reshape data between layers. A batch of images arrives as one shape, gets flattened to another, and feeds into a linear layer expecting a third. Fluency with .view() and indexing is how you control that flow.

  • Create a new code cell and explore reshaping and slicing by pasting this code:
# Cell 2 (continued): Reshaping and Indexing
a = torch.arange(12)
print(f"Original: {a}, shape={a.shape}")

# Reshape into a 3x4 matrix
reshaped = a.view(3, 4)
print(f"\nReshaped to (3,4):\n{reshaped}")

# Indexing and slicing
print(f"\nRow 0: {reshaped[0]}")
print(f"Column 1: {reshaped[:, 1]}")
print(f"Element [1,2]: {reshaped[1, 2]}")

# Add and remove dimensions
row_vec = torch.tensor([1.0, 2.0, 3.0]).unsqueeze(0)
col_vec = torch.tensor([1.0, 2.0, 3.0]).unsqueeze(1)
print(f"\nRow vector (unsqueeze 0): shape={row_vec.shape}")
print(f"Col vector (unsqueeze 1): shape={col_vec.shape}")

What does this code do?

torch.arange(12) creates a 1D tensor [0, 1, 2, ..., 11]. Calling .view(3, 4) reshapes it into 3 rows and 4 columns without copying data. The total number of elements must stay the same (3 × 4 = 12).

Indexing works like NumPy: reshaped[0] grabs row 0, reshaped[:, 1] grabs column 1 (all rows, column index 1). The .unsqueeze() method adds a dimension at the given position, turning a shape-[3] vector into either [1, 3] (a row) or [3, 1] (a column). You will need this when preparing inputs for matrix multiplication.

  • Run the cell.
  • You should see the 12-element tensor reshaped into a 3×4 grid, with individual rows, columns, and elements extracted. The unsqueeze section shows how a shape-[3] tensor becomes either [1, 3] or [3, 1].

Seeing a RuntimeError about shape?

If you see shape '[3, 4]' is invalid for input of size X, check that torch.arange(12) has exactly 12 as the argument. The product of your view dimensions (3 × 4) must equal the number of elements in the original tensor.

Help me fix a shape mismatch when reshaping tensors.

Perform matrix math with the @ operator

The single most important operation in neural networks is matrix multiplication. PyTorch uses the @ operator for it. Understanding the difference between * (element-wise) and @ (matrix multiply) is critical because nn.Linear uses @ internally.

  • Create a new code cell and explore matrix operations by pasting this code:
# Cell 2 (continued): Matrix Math
A = torch.tensor([[1.0, 2.0],
                  [3.0, 4.0]])
B = torch.tensor([[5.0, 6.0],
                  [7.0, 8.0]])

print("=== Element-wise vs Matrix Multiplication ===")
print(f"A * B (element-wise):\n{A * B}")
print(f"\nA @ B (matrix multiply):\n{A @ B}")

# Broadcasting: scalar and vector operations
print(f"\n=== Broadcasting ===")
print(f"2 * A:\n{2 * A}")

v = torch.tensor([10.0, 20.0])
print(f"\nA + v (broadcast v across rows):\n{A + v}")
print(f"A shape: {A.shape}, v shape: {v.shape}")

What does this code do?

A * B multiplies matching elements: position [0,0] times [0,0], position [0,1] times [0,1], and so on. A @ B performs true matrix multiplication: each element in the result is a dot product of a row from A and a column from B. This is what nn.Linear uses for its core computation.

Broadcasting lets PyTorch apply operations between tensors of different shapes. Adding a shape-[2] vector to a shape-[2, 2] matrix automatically adds the vector to each row. This is how nn.Linear adds its bias to every sample in a batch.

Before you run this, mentally compute A @ B at position [0,0]: it is (1×5) + (2×7) = 19. Do you expect to see 19 in the top-left of the matrix multiply output?

  • Run the cell.
  • You should see two different 2×2 results. Element-wise gives [[5, 12], [21, 32]]. Matrix multiply gives [[19, 22], [43, 50]]. The broadcasting section shows v added to every row of A.

Seeing a RuntimeError about mat1 and mat2 shapes?

Matrix multiplication requires the inner dimensions to match: for A @ B, the number of columns in A must equal the number of rows in B. If you modified the tensor sizes, check that both A and B are shape [2, 2].

Help me understand a matrix shape mismatch error.

💡 Why does this matter for nn.Linear?

Everything you just practised is exactly what happens inside nn.Linear. When you call layer(x), PyTorch runs x @ weight.T + bias. That is a matrix multiply (@), a transpose (.T), and a broadcast add (+ bias). No other operations are involved.

✔️ Awesome, I've got everything!

Great. Make sure your tensor exploration cells are saved in your notebook.

ⓧ I'd like to double check the full code

Your tensor exploration cells should contain these three sections in order. Scroll up and confirm each one is present:

  • Tensor creation: torch.tensor(), torch.zeros(), torch.randn() with shape, dtype, and device prints, plus the dtype casting section.
  • Reshaping: torch.arange(12), .view(3, 4), row/column/element indexing, and .unsqueeze() examples.
  • Matrix math: A * B element-wise, A @ B matrix multiply, scalar multiply, and vector broadcasting with A + v.

You now have hands-on fluency with tensors, shapes, dtypes, and matrix math. Next up, you will use these exact operations to crack open nn.Linear and prove there is no magic inside it.

Hand-Verify nn.Linear with Raw Math

You now know how to create tensors, reshape them, and perform matrix multiplication with the @ operator. Now it is time to see those exact operations in action inside the single most important building block in neural networks: nn.Linear.

Most people treat nn.Linear as a black box that magically transforms inputs into outputs. In this step, you will manually compute y = xW^T + b with real numbers and prove to yourself that there is no magic.

In this step, get ready to:
  • Create an nn.Linear layer and inspect its weight and bias matrices.
  • Manually replicate nn.Linear's output using raw matrix multiplication.
  • Explore how nn.Module tracks learnable parameters.
Create an nn.Linear layer and inspect its internals

An nn.Linear layer is just a container for a weight matrix and a bias vector. To verify what it does, you first need to see those values.

  • Create a new code cell below your setup cell. Add the following code to create a tiny layer and print its internals:
# Cell 2: Hand-Verify nn.Linear
torch.manual_seed(42)
layer = nn.Linear(3, 2)

print("=== nn.Linear Internals ===")
print(f"Weight shape: {layer.weight.shape}")
print(f"Weight:\n{layer.weight}")
print(f"Bias shape: {layer.bias.shape}")
print(f"Bias: {layer.bias}")

What does this code do?

torch.manual_seed(42) resets the random seed so the layer's weights are the same every time you run this cell.

nn.Linear(3, 2) creates a layer with 3 inputs and 2 outputs. The weight matrix has shape (out_features, in_features), which is [2, 3]. This is because the formula is y = xW^T + b. The transpose means PyTorch stores the weight as (output size, input size).

  • Run the cell.
  • You should see Weight shape: torch.Size([2, 3]) and Bias shape: torch.Size([2]), followed by the actual numeric values of both.

Seeing a NameError?

Make sure you ran Cell 1 (your setup cell) first. The nn import from import torch.nn as nn must execute before this cell can use it.

Help me figure out why my nn.Linear cell throws a NameError in Colab.

Manually replicate nn.Linear's output

Now that you can see the exact weight and bias values, you can do the same math PyTorch does and prove the results match.

  • Add the following code below your existing code in Cell 2:
x = torch.tensor([[1.0, 2.0, 3.0]])
output = layer(x)
print(f"\nPyTorch output: {output}")

manual_output = x @ layer.weight.T + layer.bias
print(f"Manual output:  {manual_output}")
print(f"Match: {torch.allclose(output, manual_output)}")

What does this code do?

layer(x) passes the input through PyTorch's built-in forward pass. Then x @ layer.weight.T + layer.bias computes the exact same thing using raw matrix multiplication.

The @ operator performs matrix multiplication. layer.weight.T transposes the weight from shape [2, 3] to [3, 2], so the multiplication [1, 3] @ [3, 2] produces [1, 2]. Adding the bias [2] gives the final output.

Before you run this, what do you think the Match line will print? Will your manual math produce the exact same numbers as PyTorch?

  • Run the cell.
  • You should see both outputs display the same tensor values, and Match: True at the bottom. That single True proves nn.Linear is nothing more than y = xW^T + b.

Seeing Match: False?

Check that you used layer.weight.T (with the capital T for transpose), not just layer.weight. Without the transpose, the shapes won't align for multiplication.

Help me debug why my manual output doesn't match nn.Linear's output.

Explore how nn.Module tracks parameters

You have proved that nn.Linear is just matrix math. So why use it instead of raw @ operations? The answer is nn.Module's parameter tracking system.

  • Add the following code at the bottom of Cell 2:
print("\n=== nn.Module Parameter Tracking ===")
for name, param in layer.named_parameters():
    print(f"  {name}: shape={param.shape}, requires_grad={param.requires_grad}")

What does this code do?

named_parameters() iterates over every learnable tensor the layer holds. Each weight and bias is stored as an nn.Parameter, which is a special Tensor with requires_grad=True.

This is how PyTorch knows which tensors to update during training. When you call an optimizer's step() method later, it finds these parameters automatically.

  • Run the cell.
  • You should see two entries: weight with shape torch.Size([2, 3]) and bias with shape torch.Size([2]). Both show requires_grad=True.

Not seeing requires_grad=True?

If you created the layer outside of the nn.Linear constructor (e.g., using a plain torch.randn tensor), it won't be an nn.Parameter. Make sure you used nn.Linear(3, 2) to create the layer.

Help me understand why my layer parameters don't show requires_grad=True.

✔️ Awesome, I've got everything!

Great. Make sure Cell 2 is saved in your notebook.

ⓧ I'd like to double check the full code

Your Cell 2 should contain the following code from top to bottom. Scroll up and confirm each section is present in order:

  • The torch.manual_seed(42) and nn.Linear(3, 2) creation, followed by the five print statements showing weight and bias internals.
  • The input tensor x = torch.tensor([[1.0, 2.0, 3.0]]), PyTorch output, manual output with x @ layer.weight.T + layer.bias, and the torch.allclose match check.
  • The parameter tracking loop using layer.named_parameters() printing each parameter's name, shape, and requires_grad status.

You just proved that nn.Linear is nothing more than a matrix multiplication and a bias addition. Next up, you will visualize the activation functions that introduce nonlinearity between these linear layers.

Visualize Activation Functions and Their Derivatives

You now know exactly what nn.Linear computes under the hood. But a linear layer alone just does matrix multiplication. The real power of neural networks comes from activation functions between linear layers, which introduce nonlinearity.

The problem is that activation function behavior is invisible until you see it graphed. In this step, you will plot four major activation functions and their derivatives, revealing exactly where each one saturates, dies, or stays linear.

In this step, get ready to:
  • Compute ReLU, Sigmoid, Tanh, and LeakyReLU outputs over a range of inputs.
  • Use PyTorch autograd to compute derivatives automatically.
  • Plot each function alongside its derivative in a 4x2 grid and hand-verify key values.
Compute activations and derivatives with autograd

Each activation function transforms its input in a different way. To understand these transformations, you need both the function output AND its derivative (gradient). PyTorch's autograd system can compute derivatives automatically.

  • Create a new code cell in your Colab notebook and paste the following code to set up the activations dictionary and compute each function with its gradient:
# Cell 3: Activation Functions and Derivatives
fig, axes = plt.subplots(4, 2, figsize=(12, 16))
activations = {
    'ReLU': torch.relu,
    'Sigmoid': torch.sigmoid,
    'Tanh': torch.tanh,
    'LeakyReLU': lambda x: F.leaky_relu(x, negative_slope=0.01),
}

for i, (name, func) in enumerate(activations.items()):
    x = torch.linspace(-5, 5, 200, requires_grad=True)
    y = func(x)
    y.backward(torch.ones_like(x))
    grad = x.grad.clone()

What does this code do?

The activations dictionary maps four function names to their PyTorch implementations. For each one, the loop creates 200 evenly spaced inputs from -5 to 5.

Setting requires_grad=True tells PyTorch to track operations on x. When you call .backward(torch.ones_like(x)), PyTorch computes the derivative of the function at every point and stores it in x.grad. This is the same autograd engine that powers backpropagation during training.

Plot each function alongside its derivative
  • Continue in the same cell by pasting this plotting code directly below (still inside the for loop, indented with 4 spaces):
    axes[i, 0].plot(x.detach().numpy(), y.detach().numpy(), 'b-', linewidth=2)
    axes[i, 0].set_title(f'{name}')
    axes[i, 0].axhline(y=0, color='k', linewidth=0.5)
    axes[i, 0].axvline(x=0, color='k', linewidth=0.5)
    axes[i, 0].grid(True, alpha=0.3)
    axes[i, 0].set_xlabel('Input')
    axes[i, 0].set_ylabel('Output')

    axes[i, 1].plot(x.detach().numpy(), grad.numpy(), 'r-', linewidth=2)
    axes[i, 1].set_title(f'{name} Derivative')
    axes[i, 1].axhline(y=0, color='k', linewidth=0.5)
    axes[i, 1].axvline(x=0, color='k', linewidth=0.5)
    axes[i, 1].grid(True, alpha=0.3)
    axes[i, 1].set_xlabel('Input')
    axes[i, 1].set_ylabel('Gradient')

plt.tight_layout()
plt.savefig('activation_functions.png', dpi=100)
plt.show()

What does this plotting code do?

For each activation function, the left column plots the function itself (blue) and the right column plots its derivative (red). The .detach().numpy() calls convert PyTorch tensors back to NumPy arrays for matplotlib.

The three lines outside the loop (plt.tight_layout(), plt.savefig(), plt.show()) arrange the subplots neatly, save the grid as activation_functions.png, and display it inline.

Before you run this cell, what do you think the ReLU derivative will look like for negative inputs?

  • Run the cell by pressing Shift+Enter.

You should see a 4x2 grid of plots with each activation function in blue on the left and its derivative in red on the right.

What should I notice in these plots?

  • ReLU: the derivative is exactly 0 for all negative inputs. This is the "dead zone" where neurons stop learning entirely.
  • Sigmoid: the derivative drops below 0.1 for inputs beyond roughly -2 and +2. This is the "saturation zone" where gradients vanish and training stalls.
  • Tanh: similar shape to Sigmoid, but the output is zero-centered (ranges from -1 to 1 instead of 0 to 1). Zero-centering helps subsequent layers converge faster.
  • LeakyReLU: instead of a flat zero for negatives, it has a tiny slope (0.01). This prevents dead neurons by always allowing some gradient to flow.

🙋‍♀️ Seeing an empty plot or error?

  • Check that the plotting code is indented inside the for loop (4 spaces). Only plt.tight_layout(), plt.savefig(), and plt.show() should be at the top level.
  • If you see a RuntimeError about gradients, make sure requires_grad=True is inside the torch.linspace call, not on a separate line.

Help me debug my activation function plots

Hand-verify Sigmoid at x = 0

Plots give intuition, but you should always verify numbers you can calculate by hand. Sigmoid(0) has a known closed-form answer.

  • Add this code to the bottom of the same cell, below plt.show():
# Hand-verify Sigmoid at 0
print(f"Sigmoid(0) = {torch.sigmoid(torch.tensor(0.0)):.4f} (expected 0.5)")
sig_x = torch.tensor(0.0, requires_grad=True)
sig_y = torch.sigmoid(sig_x)
sig_y.backward()
print(f"Sigmoid'(0) = {sig_x.grad:.4f} (expected 0.25)")

Why these specific values?

Sigmoid is defined as 1/(1 + e^(-x)). At x=0, that gives 1/(1+1) = 0.5.

The derivative of Sigmoid is sigmoid(x) * (1 - sigmoid(x)). At x=0, that gives 0.5 * 0.5 = 0.25. If PyTorch returns these exact values, autograd is computing derivatives correctly.

What do you expect the two printed values to be?

  • Re-run the cell with Shift+Enter.

Below your plots, you should see two printed lines confirming Sigmoid(0) = 0.5000 and Sigmoid'(0) = 0.2500.

Values don't match expected?

  • Confirm you used torch.tensor(0.0) (a float) not torch.tensor(0) (an integer). Integer tensors cannot track gradients.
  • Make sure sig_y.backward() is called before reading sig_x.grad. The gradient only exists after backward runs.

Help me debug Sigmoid verification

✔️ Awesome, I've got everything!

Great. Make sure your cell is saved. You now have a complete visual reference for the four most common activation functions.

ⓧ I'd like to double check the full code

Your Cell 3 should contain all of the following sections in order:

  • The figure and activations dictionary setup (fig, axes = plt.subplots(4, 2, ...) through the activations dict).
  • The for loop computing y and grad using backward and x.grad.clone().
  • The plotting section inside the loop (axes[i, 0] and axes[i, 1] plots with titles, axis lines, and grid).
  • The display calls at the top level: plt.tight_layout(), plt.savefig('activation_functions.png', dpi=100), plt.show().
  • The hand-verify section printing Sigmoid(0) and Sigmoid'(0).

You now have a visual atlas of activation function behavior. Next up, you will trace a complete forward pass through a 2-layer network, watching how a linear layer and ReLU work together to transform data step by step.

Trace a Full Forward Pass Through a 2-Layer Network

In the previous step, you plotted every major activation function and saw exactly where each one saturates or dies.

A single nn.Linear layer is just matrix multiplication. The real power of neural networks comes from stacking layers with activations between them.

In this step, you will build a 2-layer network and trace every intermediate value through it. By the end, you will have proof that forward() is just chained matrix multiplications and activations.

In this step, get ready to:
  • Build a 2-layer network and inspect its architecture and parameters.
  • Pass an input through layer by layer, verifying each intermediate tensor manually.
  • Confirm the full forward pass matches your hand calculations.
Build the 2-layer network

nn.Sequential is a PyTorch container that chains modules in order. During the forward pass, it feeds the output of each layer as the input to the next.

  • Create a new cell in your Colab notebook. Build a 2-layer network and print its architecture by pasting this code:
# Cell 4: Full Forward Pass Trace
torch.manual_seed(0)
net = nn.Sequential(
    nn.Linear(2, 3),
    nn.ReLU(),
    nn.Linear(3, 1)
)

print("=== Network Architecture ===")
print(net)
print("\n=== All Parameters ===")
for name, param in net.named_parameters():
    print(f"  {name}: shape={param.shape}")

What does this code do?

torch.manual_seed(0) fixes the random initialization so your weights match this walkthrough exactly.

The network has three components: a linear layer that maps 2 inputs to 3 hidden values, a ReLU activation that zeros out negatives, and a second linear layer that maps 3 hidden values down to 1 output.

named_parameters() lists every learnable weight and bias that PyTorch is tracking. You will use these exact values for hand-verification next.

  • Run the cell.

You should see the network architecture printed as a numbered list of modules, followed by four parameters: 0.weight (shape [3, 2]), 0.bias (shape [3]), 2.weight (shape [1, 3]), and 2.bias (shape [1]).

Not seeing the output?

Make sure you ran the Setup cell (Cell 1) first in this session. The torch and nn imports need to be active.

Help me debug why my nn.Sequential network isn't printing

Trace the first linear layer manually

Now you will feed an input through the first layer and verify the output matches the formula y = xW^T + b using the exact weight values PyTorch generated.

  • Add this code to the same cell, below what you just wrote:
x = torch.tensor([[1.0, -1.0]])
print(f"\nInput: {x}")

h = net[0](x)
print(f"After Linear(2,3): {h}")

manual_h = x @ net[0].weight.T + net[0].bias
print(f"Manual verification: {manual_h}")
print(f"Match: {torch.allclose(h, manual_h)}")

What does this code do?

net[0] accesses the first module in your Sequential (the nn.Linear(2, 3) layer). Calling net[0](x) passes x through just that one layer.

The manual verification repeats the same x @ weight.T + bias formula you proved in Step 2. If both outputs match, you have confirmed there is no hidden magic between layers.

  • Run the cell.

You should see your input tensor, the hidden layer output h with 3 values, and Match: True confirming your manual calculation is identical to PyTorch's.

Seeing Match: False?

Check that torch.manual_seed(0) appears before the nn.Sequential line. If the seed is missing or set to a different number, the weights will differ from what the rest of the cell expects.

Why is my manual matrix multiplication not matching PyTorch's output?

Verify ReLU and the final output

The hidden values h now pass through ReLU (any negative becomes 0) and then through the second linear layer to produce the final scalar output.

  • Add this code to the same cell to complete the forward pass trace:
a = net[1](h)
print(f"After ReLU: {a}")

manual_a = torch.relu(h)
print(f"Manual ReLU: {manual_a}")

out = net[2](a)
print(f"After Linear(3,1): {out}")

manual_out = a @ net[2].weight.T + net[2].bias
print(f"Manual verification: {manual_out}")
print(f"Match: {torch.allclose(out, manual_out)}")

full_out = net(x)
print(f"\nFull forward pass: {full_out}")
print(f"Matches manual trace: {torch.allclose(full_out, out)}")

What does this code do?

net[1](h) applies ReLU to the hidden values. You verify this by calling torch.relu(h) directly. Any negative value in h becomes 0.

net[2](a) runs the second linear layer. You verify with the same a @ weight.T + bias formula.

Finally, net(x) runs the entire forward pass in one call. If it matches your manually traced out tensor, you have proven that forward() just chains these three computations.

Before you run the cell, think about this: if any value in h is negative, what will happen to it after ReLU? And will the final net(x) output match your layer-by-layer trace?

  • Run the cell.

You should see every intermediate tensor printed: the ReLU output (negatives become 0), the final linear output, and Matches manual trace: True at the bottom.

What just happened?

You traced a complete forward pass through a 2-layer network by hand. Every intermediate value matched your manual calculation.

forward() is not magic. It is matrix multiplication, then activation, then matrix multiplication again. That is the entire computation graph of a feed-forward network.

🙋‍♀️ Matches manual trace shows False?

Make sure you did not re-run just part of the cell. If you re-ran the top half (which recreates net with new random weights) without re-running the bottom half, the variables h and a still hold old values. Run the entire cell from top to bottom.

Why doesn't my full net(x) forward pass match my manual trace?

✔️ Awesome, I've got everything!

Double check you have all the code in Cell 4 saved. Your cell should contain the network creation, the layer-by-layer trace, and the full forward pass verification.

ⓧ I'd like to double check the full code

Your Cell 4 should contain these three sections in order. Scroll up and confirm each one is present:

  • Network creation: torch.manual_seed(0), nn.Sequential definition, architecture print, and parameters loop.
  • First layer trace: input x creation, net[0](x), manual_h verification with allclose.
  • ReLU and final output: net[1](h), manual_a, net[2](a), manual_out, and the full net(x) comparison.

You have now traced every number through a multi-layer network and confirmed there is no hidden step. You understand exactly what happens when PyTorch calls forward(). Next up is the Secret Mission: building your own Linear layer from scratch.

Secret mission

Build Your Own Linear Layer

You have verified what nn.Linear does from the outside. Now prove you understand it from the inside by reimplementing it from scratch using nn.Parameter, and showing PyTorch cannot tell the difference.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely on Google Colab, so there are no ongoing costs.

Resources you used:

  • Google Colab notebook (stored in your Google Drive).
  • Colab runtime session (GPU/CPU compute time).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with your notebook or want to revisit the activation function plots later.

  • Your notebook stays in your Google Drive under the Colab Notebooks folder.
  • The runtime will automatically disconnect after a period of inactivity, so there is nothing to manage.

Pause - I'll come back to this later

Disconnect your runtime to free up shared compute resources. Your notebook code stays saved in Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • This frees up the GPU/CPU session. Any files saved only in the runtime (like activation_functions.png) will be lost, but your notebook code remains in Google Drive.
  • When you return, open the notebook from Google Drive and click Runtime then Run all to regenerate your outputs.

Delete - I don't want to use this again

Remove the notebook from your Google Drive entirely.

  • In your Colab notebook, click Runtime in the top menu bar, then Disconnect and delete runtime.
  • Open Google Drive at drive.google.com.
  • Navigate to the Colab Notebooks folder.
  • Right-click the notebook file and select Move to trash.
  • Empty your trash to permanently delete it: click Trash in the left sidebar, then click Empty trash.

Nice Work!

Nice Work!

Nice work! You just built a complete PyTorch nn internals reference notebook from scratch, hand-verifying every number along the way.

You've learned how to:

  • Hand-verify that nn.Linear computes y = xW^T + b by extracting weights and replicating the output with raw matrix multiplication.
  • Visualize activation functions (ReLU, Sigmoid, Tanh, LeakyReLU) alongside their derivatives, identifying saturation zones and dead zones directly from the plots.
  • Trace a complete forward pass through a 2-layer network, verifying every intermediate tensor against your own calculations using nn.Module internals like named_parameters() and requires_grad.
  • Secret Mission: Built your own custom Linear layer from scratch using nn.Parameter and proved it produces identical results to nn.Linear.

Ready to quiz yourself?