Machine Learning Academy · 课时

使用 nn.Module 构建前馈网络

您将继承 nn.Module，在 __init__ 中堆叠 Linear 和 ReLU 层，实现前向传播，并使用虚拟输入验证输出形状。

第 3 / 4 课13 个步骤

使用 nn.Module 构建前馈网络 是 CoddyKit 上的免费 Machine Learning Academy 课时。这是第 3 节课，共 4 节。你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。这是 Machine Learning Academy 学习路径的一部分，你的进度在网页和 CoddyKit 应用中同步。 Machine Learning Academy 课程共包含 4 节课。

本课时的部分内容尚未翻译，以英文显示。

What Is nn.Module?

nn.Module is PyTorch's base class for all neural network components. Every layer, activation function, loss function, and complete model in PyTorch subclasses nn.Module. It provides parameter management, device movement, serialisation, and training/eval mode switching out of the box. By subclassing it you get all this functionality for free and only need to define the architecture and the forward pass.

import torch.nn as nn

# Inspect what nn.Module gives you
model = nn.Linear(4, 2)
print(type(model))               # <class 'torch.nn.modules.linear.Linear'>
print(isinstance(model, nn.Module))  # True
print(list(model.parameters()))  # weight and bias tensors

Subclassing nn.Module: init and forward

Creating a custom network requires two methods. In __init__ you call super().__init__() and define all learnable layers as attributes. In forward you describe how data flows through those layers. PyTorch tracks any nn.Module or nn.Parameter assigned to self as a trainable component. Calling the model like a function (model(x)) automatically invokes forward and runs any registered hooks.

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet(4, 16, 3)
print(model)

Linear Layers: nn.Linear

nn.Linear(in_features, out_features) applies the transformation y = x @ W.T + b where W is a weight matrix and b is a bias vector. The layer automatically initialises weights with kaiming uniform distribution and biases with uniform distribution. It is the building block of feedforward networks, also called fully connected or dense layers.

import torch
import torch.nn as nn

layer = nn.Linear(5, 3)  # 5 inputs, 3 outputs

print('Weight shape:', layer.weight.shape)  # (3, 5)
print('Bias shape:',   layer.bias.shape)    # (3,)

# Forward pass with a batch of 8 samples
x = torch.randn(8, 5)
out = layer(x)
print('Output shape:', out.shape)   # (8, 3)

Activation Functions: ReLU, Sigmoid, Tanh

Activation functions introduce non-linearity, allowing the network to learn complex patterns that a stack of linear layers cannot represent. ReLU (max(0, x)) is the default for hidden layers — fast and avoids vanishing gradients. Sigmoid squashes output to [0, 1], used for binary classification outputs. Tanh squashes to [-1, 1], commonly used in RNNs. All are available as modules in nn.

import torch
import torch.nn as nn

x = torch.tensor([-2.0, -0.5, 0.0, 0.5, 2.0])

print('ReLU:   ', nn.ReLU()(x))
# tensor([0.0, 0.0, 0.0, 0.5, 2.0])

print('Sigmoid:', nn.Sigmoid()(x))
# tensor([0.12, 0.38, 0.50, 0.62, 0.88])

print('Tanh:   ', nn.Tanh()(x))
# tensor([-0.96, -0.46,  0.00,  0.46,  0.96])

Stacking Layers with nn.Sequential

nn.Sequential is a convenient container that passes the output of each module as the input to the next. It is ideal for simple feedforward architectures where data flows linearly. For more complex networks with skip connections, multiple inputs, or branching paths, you need the full nn.Module subclass approach. Sequential networks can still be extended by subclassing and using the Sequential as a sub-block.

import torch
import torch.nn as nn

# Build with nn.Sequential
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

x = torch.randn(16, 10)    # batch of 16
out = model(x)
print(out.shape)            # torch.Size([16, 1])

Verifying Output Shapes with Dummy Input

A critical debugging technique when building networks is to run a dummy tensor through the model before training. This verifies that all layer dimensions are compatible and reveals shape mismatches immediately. The dummy tensor has the same shape as your real data but contains random values — it just exercises the forward path. Always do this after defining a new architecture.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    def forward(self, x):
        return self.net(x)

model = MLP()
dummy = torch.randn(32, 784)   # batch of 32 MNIST images
out = model(dummy)
print(out.shape)                # torch.Size([32, 10]) -- correct!

Listing and Counting Parameters

Understanding the total number of trainable parameters in a model gives you a sense of its capacity and memory requirements. model.parameters() returns an iterator over all learnable tensors; model.named_parameters() pairs each tensor with its name for inspection. A compact parameter counter is one of the first utilities every PyTorch practitioner builds.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters()
                if p.requires_grad)

print(f'Total params:     {total_params:,}')   # 203,530
print(f'Trainable params: {trainable:,}')       # 203,530

for name, p in model.named_parameters():
    print(name, p.shape)

Moving the Model to GPU

Moving the model to the GPU is as simple as calling model.to(device). This transfers all parameters and buffers registered in the module to the target device. After this call, all forward pass computations happen on the GPU automatically — as long as your input tensors are also on the same device. Mixing CPU and GPU tensors raises a runtime error.

import torch
import torch.nn as nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2)
).to(device)   # move entire model to device

# Input must be on the same device
x = torch.randn(5, 4).to(device)
out = model(x)
print(out.device)   # cuda:0 (or cpu)

Training vs Eval Mode

Some layers (Dropout, BatchNorm) behave differently during training and inference. Calling model.train() enables stochastic behaviour (random dropout, batch statistics); model.eval() switches to deterministic inference behaviour. Forgetting to call model.eval() before evaluation leads to inconsistent results because Dropout randomly zeros activations. Always pair these calls with torch.no_grad() during inference for maximum efficiency.

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.Dropout(p=0.5),
    nn.Linear(8, 2)
)

# Training mode: dropout is active
model.train()
x = torch.randn(4, 4)
print(model(x))   # some activations zeroed randomly

# Eval mode: dropout is disabled
model.eval()
with torch.no_grad():
    print(model(x))   # deterministic output

Custom Network: Multi-Layer Perceptron

Putting it all together: a multi-layer perceptron (MLP) is a feedforward network with one or more hidden layers between input and output. Each hidden layer applies a linear transformation followed by a non-linear activation. The output layer's activation depends on the task — no activation for regression, softmax for multi-class classification, sigmoid for binary classification. The example below builds a 3-layer MLP for 10-class classification.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden_dims, out_dim):
        super().__init__()
        layers = []
        prev = in_dim
        for h in hidden_dims:
            layers.append(nn.Linear(prev, h))
            layers.append(nn.ReLU())
            prev = h
        layers.append(nn.Linear(prev, out_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)

model = MLP(784, [256, 128, 64], 10)
dummy = torch.randn(32, 784)
print(model(dummy).shape)  # torch.Size([32, 10])

Saving and Loading Network State

Neural network training is expensive, so you save the model after training. PyTorch's convention is to save only the state_dict — a dictionary of parameter tensors — rather than the entire model object. This avoids pickle dependency on the class definition. To restore, create a fresh model instance with the same architecture, then load the state dict with load_state_dict. Always set model.eval() after loading for inference.

import torch
import torch.nn as nn

model = nn.Linear(4, 2)

# Save only parameters (recommended)
torch.save(model.state_dict(), '/tmp/model.pt')

# Restore
new_model = nn.Linear(4, 2)  # same architecture
new_model.load_state_dict(torch.load('/tmp/model.pt'))
new_model.eval()

x = torch.randn(3, 4)
print(new_model(x))   # same output as original model

Quick Check

Test your understanding of Machine Learning with Python concepts from this lesson.

在交互式课程中回答

Lesson Recap

In this lesson you learned: nn.Module is the base class for all PyTorch neural networks, requiring __init__ to define layers and forward to define data flow, nn.Sequential provides a simple container for linear stacks of layers, and model.train() / model.eval() switches behaviour of layers like Dropout and BatchNorm. Next up we write the complete training loop including loss, optimizer, and multiple epochs.

免费开始

用 AI 导师学习 Python — 免费

在浏览器中编写并运行真实代码，获得全天候 AI 导师的即时帮助，并在网页或应用中继续学习。

课程: 30
课程: 120

免费开始学习查看完整课程

常见问题解答