🎓 A1Learner
🔍
0 of 71 lessons completed

PyTorch A Reusable Residual Block

A residual block is a module that adds its input back to its output. The skip connection is the input carried straight to that final add.

One block, packaged once

The block holds one linear layer, one LayerNorm, and one ReLU. forward adds the input back after those three parts.

class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)
        self.norm = nn.LayerNorm(dim)
        self.act = nn.ReLU()

    def forward(self, x):
        return x + self.act(self.norm(self.linear(x)))

One block keeps the shape

One block returns the same shape it receives. The skip add needs equal shapes, so the block keeps the width fixed.

torch.manual_seed(0)
block = ResidualBlock(8)
x = torch.randn(4, 8)
out = block(x)
print("Input shape:", tuple(x.shape))
print("Output shape:", tuple(out.shape))

The same class, stacked

The same class builds every block in this stack. Three blocks hold three times one block's parameters.

stack = nn.Sequential(ResidualBlock(8), ResidualBlock(8), ResidualBlock(8))
print("Blocks:", len(stack))
print("Output shape:", tuple(stack(x).shape))
print("Parameters:", sum(p.numel() for p in stack.parameters()))

Example

import torch
from torch import nn


class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)
        self.norm = nn.LayerNorm(dim)
        self.act = nn.ReLU()

    def forward(self, x):
        return x + self.act(self.norm(self.linear(x)))


torch.manual_seed(0)
stack = nn.Sequential(*[ResidualBlock(8) for _ in range(6)])
x = torch.randn(4, 8)
out = stack(x)

total = sum(p.numel() for p in stack.parameters())
per_block = sum(p.numel() for p in stack[0].parameters())

print("Blocks:", len(stack))
print("Input shape:", tuple(x.shape))
print("Output shape:", tuple(out.shape))
print("Parameters per block:", per_block)
print("Total parameters:", total)

Key points

  • Residual block: a module that returns its input plus a transformed copy
  • Skip connection: the path that carries the input straight to the add
  • Use x + self.act(self.norm(self.linear(x))) in forward to add the skip
  • Every block in a stack shares one width, so shapes stay fixed
  • Stacking reuses one class, so depth costs no new code

In this example, six ResidualBlock objects are stacked in one nn.Sequential. The input holds four rows of eight numbers. Every block returns the same shape it receives, so the output keeps the shape (4, 8). Each block holds 88 parameters, from its linear layer and its LayerNorm. Six blocks total 528 parameters. One class definition produced all six blocks, with no copied code.

Note

LayerNorm inside the block pins each row's scale as depth grows

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why must a residual block return the same shape it receives?