PyTorch神经网络实现：从二分类到Transformer

1. 从零开始：PyTorch神经网络实现三部曲

在深度学习领域，PyTorch因其动态计算图和直观的API设计，成为众多研究者和工程师的首选框架。今天我将分享三种经典神经网络结构的PyTorch实现方法，这些代码虽然简单，但包含了模型构建的核心逻辑。无论你是刚入门的新手还是需要快速回顾的老手，这些实现都能帮助你理解神经网络的工作机制。

2. 二分类神经网络：理解基础架构

2.1 数据准备与模型定义

我们先从一个简单的二分类任务开始。在PyTorch中，数据通常以张量(Tensor)形式组织。对于这个例子，我们随机生成10个样本，每个样本有10个特征：

python复制n_in, n_h, n_out, batch_size = 10, 5, 1, 10
x = torch.randn(batch_size, n_in)  # 输入数据
y = torch.tensor([[1.0], [0.0], [0.0], [1.0], [1.0], 
                 [1.0], [0.0], [0.0], [1.0], [1.0]])  # 目标输出

这里y采用二维张量结构是为了保持与复杂任务(如多分类)的兼容性。即使输出是单个标量，也建议保持这种格式。

PyTorch提供了两种定义模型的方式：Sequential和Module类。Sequential适合简单线性结构：

python复制model = nn.Sequential(
   nn.Linear(n_in, n_h),
   nn.ReLU(),
   nn.Linear(n_h, n_out),
   nn.Sigmoid()
)

而Module类则提供了更大的灵活性：

python复制class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(n_in, n_h)
        self.fc2 = nn.Linear(n_h, n_out)
        self.sigmoid = nn.Sigmoid()

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

提示：在实际项目中，Module类是更推荐的做法，因为它允许你在forward方法中实现任意复杂的前向逻辑。

2.2 训练流程与关键细节

定义好模型后，我们需要设置损失函数和优化器：

python复制criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

训练循环包含几个关键步骤：

python复制for epoch in range(50):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    
    optimizer.zero_grad()  # 清除上一轮的梯度
    loss.backward()        # 反向传播计算梯度
    optimizer.step()       # 更新参数
    
    print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

这里有几个容易忽视但重要的细节：

1. zero_grad()必须在backward()之前调用，否则梯度会累积
2. 在某些场景下，可以累积多个batch的梯度后再更新，这对处理大batch或内存有限的情况很有帮助
3. loss.item()将单元素张量转换为Python数值，避免不必要的计算图保留

3. 卷积神经网络：图像处理利器

3.1 数据加载与预处理

对于图像任务，PyTorch提供了torchvision工具包简化数据处理。以MNIST手写数字识别为例：

python复制transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 将[0,1]范围归一化到[-1,1]
])

train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

归一化处理有助于模型训练的稳定性。这里的(0.5,), (0.5,)表示对单通道图像进行均值和标准差归一化，计算公式为：x_norm = (x - mean) / std。

3.2 CNN架构设计

CNN通过局部连接和权值共享显著减少了参数数量。一个典型的CNN结构如下：

python复制class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64*7*7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 64*7*7)  # 展平
        x = F.relu(self.fc1(x))
        return self.fc2(x)

关键组件说明：

3.3 训练技巧

CNN训练通常使用交叉熵损失和带动量的SGD：

python复制model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

model.train()  # 设置训练模式(启用dropout等)
for epoch in range(5):
    for images, labels in train_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

注意：model.train()和model.eval()的切换很重要，它会影响dropout和batchnorm等层的行为。在验证和测试时记得切换到eval模式。

4. Transformer：注意力机制的威力

4.1 位置编码

Transformer抛弃了RNN的循环结构，改用注意力机制处理序列数据。由于没有循环结构，需要显式地加入位置信息：

python复制class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        
        pe = torch.zeros(1, max_len, d_model)
        pe[0, :, 0::2] = torch.sin(position * div_term)  # 偶数位置
        pe[0, :, 1::2] = torch.cos(position * div_term)  # 奇数位置
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

这种正弦余弦编码满足：

1. 每个位置有唯一编码
2. 相对位置关系可通过简单线性变换建模
3. 可处理比训练时更长的序列

4.2 注意力机制

多头注意力是Transformer的核心：

python复制class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, Q, K, V, mask=None):
        # 线性变换并分头
        Q = self.W_q(Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        K = self.W_k(K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        V = self.W_v(V).view(batch_size, -1, n_heads, d_k).transpose(1, 2)
        
        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)  # (batch_size, n_heads, seq_len, d_k)
        
        # 合并多头输出
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, d_model)
        return self.W_o(output)

关键点解析：

1. 查询(Q)、键(K)、值(V)通过线性变换得到
2. 分数计算采用缩放点积注意力，防止梯度消失
3. 掩码机制用于处理变长序列和防止信息泄露

4.3 训练中的实用技巧

训练Transformer时需要注意：

1. 学习率预热：初始阶段逐步增大学习率，有助于稳定训练
2. 标签平滑：防止模型对标签过度自信
3. 梯度裁剪：避免梯度爆炸
4. 学习率调度：如余弦退火等策略

python复制optimizer = optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)
lr_scheduler = LambdaLR(
    optimizer,
    lr_lambda=lambda step: min((step+1)**-0.5, (step+1)*warmup_steps**-1.5)
)

for step in range(total_steps):
    optimizer.zero_grad()
    loss = model(x, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    lr_scheduler.step()

5. 常见问题与调试技巧

5.1 梯度消失/爆炸

症状：训练早期loss不下降或变为NaN
解决方法：

• 使用梯度裁剪(clip_grad_norm_)
• 合适的权重初始化(如Xavier初始化)
• 添加残差连接
• 使用Layer Normalization

5.2 过拟合

症状：训练误差低但验证误差高
解决方法：

• 增加Dropout
• 数据增强
• 权重衰减(L2正则化)
• 早停(Early Stopping)

5.3 训练不稳定

症状：loss波动大
解决方法：

• 减小学习率
• 增加batch size
• 使用学习率预热
• 检查数据预处理是否一致

5.4 模型评估指标

除了loss，还应该监控：

• 准确率(分类任务)
• 困惑度(语言模型)
• BLEU分数(机器翻译)
• 其他任务相关指标

6. 性能优化技巧

6.1 混合精度训练

python复制scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

6.2 数据加载优化

使用DataLoader的num_workers参数并行加载数据：

python复制DataLoader(..., num_workers=4, pin_memory=True)

6.3 模型并行

对于大模型：

python复制model = nn.DataParallel(model)  # 数据并行
# 或者
model = model.to('cuda:0')
part_of_model = part_of_model.to('cuda:1')  # 模型并行

6.4 内存优化

• 使用梯度检查点(checkpointing)
• 及时释放不需要的张量(del tensor)
• 使用torch.no_grad()上下文

7. 从理论到实践的思考

在实际项目中，单纯实现模型结构只是开始。有几个更深层次的考虑：

1. 可解释性：如何理解模型的决策过程？注意力权重可以提供一定洞察
2. 部署考量：模型需要满足推理延迟和内存占用的约束
3. 领域适配：预训练+微调范式在NLP中很成功，但在其他领域如何应用？
4. 计算效率：Transformer的O(n²)复杂度如何优化？稀疏注意力、分块处理等策略

我个人的经验是，理解底层原理比单纯调用高级API更重要。当出现问题时，扎实的基础知识能帮助你快速定位和解决。例如，当注意力分数全部接近零时，知道检查缩放因子和初始化方式就能节省大量调试时间。