PyTorch

Module

The parent class of all models; every network inherits from it

  • init(self)
  • forward(self, x)

conv

Bilibili

1
output=F.conv2d(input, kernel, stride=1) # performs the convolution computation, producing the feature matrix

Conv2d

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(
            in_channels=3,
            out_channels=6,
            kernel_size=3,
            stride=1,
            padding=0
        )

        self.maxpool1 = nn.MaxPool2d(
            kernel_size=3,  # Pooling window size: 3×3
            ceil_mode=True  # Ceiling mode, keeps incomplete edge windows
            # stride defaults to kernel_size=3; you can also explicitly write stride=3, with exactly the same effect
        )

    def forward(self, x):
        # Forward pass: **first convolve to extract features, then pool to downsample** (standard CNN order)
        # x: [b, 3, h, w]
        x = self.conv1(x)       # Step 1: pass through the conv layer, producing a 6-channel feature map, [b, 6, h', w']
        # h'=(h+2*p-k)/s+1
        x = self.maxpool1(x)    # Step 2: pass through the pooling layer, downsampling the feature map
        return x

Linear

1
2
3
4
5
self.linear = nn.Linear(
    in_features=1024,
    out_features=512,
    bias=True        # Default True; whether to add the bias b
)

weight = (512, 1024) bias = (512) y = x @ W.T + b

nn.BatchNorm2d

  • For the same channel C, all H×W pixels of every image in the entire batch are grouped together to compute the mean and variance.
  • Conv → BN → Activation
  • The parameter is feature_num
1
2
3
4
5
6
7
8
9
# Convolution block (standard structure)
self.conv1 = nn.Conv2d(3, 16, 3) #out_channel=16
self.bn1   = nn.BatchNorm2d(16)
self.relu  = nn.ReLU()

# Forward pass
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)

nn.Sequential

1
2
3
4
5
6
7
8
from collections import OrderedDict

model = nn.Sequential(OrderedDict([
    ('conv1', nn.Conv2d(1, 20, 5)),  # Name the 1st layer "conv1"
    ('relu1', nn.ReLU()),            # Name the 2nd layer "relu1"
    ('conv2', nn.Conv2d(20, 64, 5)), # Name the 3rd layer "conv2"
    ('relu2', nn.ReLU())             # Name the 4th layer "relu2"
]))

A simple way to chain these layers together

nn.CrossEntropyLoss

  • Softmax (converts the output into 0~1 probabilities)
  • Take the log
  • Cross-entropy computation
  • If a CNN is used for classification, you must flatten
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3,6,3)
        self.pool = nn.MaxPool2d(2)
        self.flatten = nn.Flatten()

        # After convolution and pooling, each sample has shape [6, 15, 15].
        self.fc = nn.Linear(6*15*15, 10)

    def forward(self, x):
        # Assume the input is [B,3,32,32]
        x = self.conv1(x)   # [B,6,30,30]  ✅ you got it right
        x = self.pool(x)    # [B,6,15,15]    ✅ you got it right
        x = self.flatten(x) # [B,6*15*15]    ✅
        x = self.fc(x)      # [B,10]
        return x
# Training
model = CNN()
criterion = nn.CrossEntropyLoss()  # dedicated to classification

for inputs, labels in dataloader:
    outputs = model(inputs)
    loss = criterion(outputs, labels)  # compute the loss
    loss.backward()

nn.MSELoss

nn.Dropout

nn.Embedding

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
# 1. Word vocabulary: word → numeric ID
vocab = {"I":1, "love":2, "AI":3, "ML":4}

# 2. Real sentences → converted to numbers
sentences = [[1,2,3], [1,2,4]]  # 2 sentences, 3 words each
input = torch.tensor(sentences)

# 3. Define the Embedding: the vocabulary has 5 words total (0-4), each word converted to a 4-dim vector
emb = nn.Embedding(num_embeddings=5, embedding_dim=4)

# 4. Look up the table: number → vector
output = emb(input)

# Result
print("Input shape:", input.shape)  # (2, 3)
print("Output shape:", output.shape)# (2, 3, 4)

nn.Softmax

1
2
3
4
5
6
7
8
# Model output: 2 images, 10 classes (raw scores, can be positive or negative)
logits = torch.randn(2, 10)

# Softmax converts to probabilities
prob = nn.Softmax(dim=-1)(logits) # applies softmax over the last dimension of the input

print(prob.shape)  # [2, 10]  shape unchanged
print(prob.sum(dim=-1))  # [1,1]  each row sums to 1

nn.MultiheadAttention

1
2
3
4
5
6
attn = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
# Input: 2 samples, each with 3 vectors, each vector 8-dim
x = torch.randn(2, 3, 8)
# Self-attention (Q=K=V use the same input)
output, _ = attn(x, x, x)
print(output.shape)  # [2, 3, 8] → the shape is exactly the same as the input

nn.Parameter

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
import torch
import torch.nn as nn

class MyLayer(nn.Module):
    def __init__(self):
        super().__init__()
        # ✅ Correct: a trainable parameter
        self.weight = nn.Parameter(torch.randn(5, 3))
        # ❌ Wrong: a plain tensor, not involved in training
        self.bias = torch.randn(5)

model = MyLayer()
# Only the weight above will be printed; bias will not appear
print(list(model.parameters()))

nn.TransformerEncoderLayer nn.Upsample nn.AdaptiveAvgPool2d nn.ConvTranspose2d nn.GELU