PyTorch 神经网络基础

nn.Module：模型的基类

继承它就能定义网络：

import torch
import torch.nn as nn

class TwoLayerNet(nn.Module):
    def __init__(self, in_dim, hidden, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)

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


model = TwoLayerNet(10, 32, 2)
x = torch.randn(4, 10)             # batch=4, 维度=10
out = model(x)                       # 自动调 forward
print(out.shape)                     # torch.Size([4, 2])

注意：调 model(x) 不要调 model.forward(x)——前者会做钩子等额外工作。

看模型有多少参数

total = sum(p.numel() for p in model.parameters())
print(f"参数总数: {total}")

# 显式列出
for name, p in model.named_parameters():
    print(name, p.shape, p.requires_grad)

常用层

nn.Linear(in_dim, out_dim)        # 全连接
nn.Conv2d(in_ch, out_ch, kernel)  # 2D 卷积
nn.LSTM(in_dim, hidden)           # LSTM
nn.Embedding(vocab, dim)          # 嵌入层（NLP 用）
nn.Dropout(p=0.5)                  # 正则化
nn.BatchNorm1d(dim)                # 批量归一化
nn.LayerNorm(dim)                  # 层归一化（Transformer 用）

激活函数

torch.relu(x)
torch.sigmoid(x)
torch.tanh(x)
torch.softmax(x, dim=-1)

# 或层版本
nn.ReLU()
nn.GELU()
nn.LeakyReLU(0.01)

Sequential：串联多个层

model = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 32),
    nn.ReLU(),
    nn.Linear(32, 2),
)

简单网络用 Sequential 一行写完，复杂网络用类继承。

损失函数

loss_fn = nn.CrossEntropyLoss()    # 分类
loss_fn = nn.MSELoss()             # 回归
loss_fn = nn.BCELoss()             # 二分类（带 sigmoid 用 BCEWithLogitsLoss）

# 用法
loss = loss_fn(predictions, targets)
loss.backward()

优化器

import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01)
optimizer = optim.Adam(model.parameters(), lr=1e-3)         # 万能默认
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)  # 现代默认

# 用法
optimizer.zero_grad()        # 清零
loss.backward()              # 算梯度
optimizer.step()             # 更新参数

模型保存 / 加载

# 保存权重
torch.save(model.state_dict(), "model.pth")

# 加载（注意：要先创建相同结构的模型）
model = TwoLayerNet(10, 32, 2)
model.load_state_dict(torch.load("model.pth"))
model.eval()                 # 切换到推理模式（关闭 dropout 等）

train() vs eval()

model.train()    # 训练模式：dropout 生效、BatchNorm 用 batch 统计
model.eval()     # 推理模式：dropout 关闭、BatchNorm 用累计统计

推理前一定要 .eval()——忘了就会出现"训练效果好、推理差"的诡异 bug。

下一篇讲完整的训练循环。