iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

Layers

By · Reviewed & last updated · Part of PyTorch Tutorial

torch.nn ships every layer you’d use day-to-day — linear, convolution, normalisation, pooling, embedding, attention. They’re Modules: track parameters, move with .to(device), save with state_dict().

A cheat-sheet of the layers you reach for

EXAMPLE
import torch.nn as nn

# 1) Linear (fully connected)
nn.Linear(in_features=128, out_features=64)

# 2) Convolutions
nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3, padding=1)   # 1D — sequences
nn.Conv2d(3, 32, kernel_size=3, padding=1)                              # 2D — images
nn.Conv3d(1, 8, kernel_size=3, padding=1)                               # 3D — volumetric
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)         # upsample (GANs / segmentation)

# 3) Recurrent / sequence
nn.LSTM(input_size=64, hidden_size=128, num_layers=2, batch_first=True, dropout=0.2)
nn.GRU(input_size=64, hidden_size=128, batch_first=True)
nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048, batch_first=True)

# 4) Normalisation
nn.BatchNorm2d(num_features=64)
nn.LayerNorm(normalized_shape=128)         # used in transformers
nn.GroupNorm(num_groups=8, num_channels=64)

# 5) Activations
nn.ReLU()  nn.LeakyReLU(0.1)  nn.GELU()  nn.SiLU()  nn.Softmax(dim=-1)

# 6) Dropout
nn.Dropout(p=0.2)
nn.Dropout2d(p=0.2)               # whole channels

# 7) Pooling
nn.MaxPool2d(kernel_size=2)
nn.AdaptiveAvgPool2d(output_size=1)        # always returns 1x1 — perfect before a Linear

# 8) Embeddings + attention
nn.Embedding(num_embeddings=10_000, embedding_dim=64, padding_idx=0)
nn.MultiheadAttention(embed_dim=512, num_heads=8, batch_first=True)

# Sequential — chain layers when no branching needed
import torch.nn as nn
import torch.nn.functional as F

net = nn.Sequential(
    nn.Conv2d(3, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
    nn.Flatten(), nn.Linear(64, 10),
)

Why it matters

Use nn.AdaptiveAvgPool2d instead of hard-coded sizes before the final Linear — input sizes can change (training crops vs eval), and the rest of the model stays the same.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
import torch.nn as nn
nn.Linear(in_features=128, out_features=64)
nn.Conv2d(3, 16, kernel_size=3, padding=1)
nn.LSTM(input_size=64, hidden_size=128, batch_first=True)
nn.LayerNorm(128)
Try it Yourself »

Discussion

Loading…