从LeNet走到Transformer:神经网络架构设计笔记

后来发现 Transformer 更通用,但迁移过程踩了不少坑。

这次做文本分类项目,从 CNN 迁移到 Transformer,踩了一些坑,也学到了不少。

从 CNN 开始

最初用的是 CNN,因为做的是文本分类,一维 CNN 效果还不错:

import torch
import torch.nn as nn

class TextCNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.conv = nn.Conv1d(embedding_dim, 100, kernel_size=3)
        self.fc = nn.Linear(100, num_classes)

    def forward(self, x):
        x = self.embedding(x)  # [batch, seq_len, embed_dim]
        x = x.transpose(1, 2)  # [batch, embed_dim, seq_len]
        x = self.conv(x)
        x = torch.relu(x)
        x = torch.max_pool1d(x, x.size(2))
        x = x.squeeze(2)
        x = self.fc(x)
        return x

CNN 的优势很明显:

  • 训练速度快
  • 计算资源要求低
  • 适合短文本分类

但问题也暴露了:

  • 长文本效果差
  • 上下文理解能力有限
  • 对位置信息不敏感

尝试 Transformer

后来决定迁移到 Transformer,主要原因是:

  • 需要处理长文本
  • 上下文理解更重要
  • 想统一架构(文本都用 Transformer)

基础 Transformer 实现

class TransformerClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, num_heads, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.positional_encoding = PositionalEncoding(embedding_dim)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=num_heads
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=6)
        self.fc = nn.Linear(embedding_dim, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        x = self.positional_encoding(x)
        x = self.transformer(x)
        x = x.mean(dim=1)  # 全局平均池化
        x = self.fc(x)
        return x

踩坑一:位置编码

一开始没加位置编码,发现 Transformer 效果比 CNN 还差。

查了半天才发现,Transformer 是位置无关的,不显式编码位置信息的话,就不知道词之间的相对位置。

解决:加上位置编码。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        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(max_len, 1, 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(0)]

踩坑二:计算资源消耗

Transformer 的计算消耗比 CNN 大太多了。

Batch size 爆显存:同样的机器,CNN 可以跑 256 的 batch,Transformer 只能跑 16。

graph TB subgraph 资源消耗对比 A[Transformer] B[CNN] A --> A1[显存: 16GB<br/>Batch: 16] A --> A2[训练时间: 2小时] B --> B1[显存: 8GB<br/>Batch: 256] B --> B2[训练时间: 20分钟] end style A1 fill:#FFB6C1,stroke:#FF0000,stroke-width:2px style A2 fill:#FFB6C1,stroke:#FF0000,stroke-width:2px style B1 fill:#90EE90 style B2 fill:#90EE90

解决

  • 用混合精度训练
  • 减少 Transformer 层数
  • 用梯度累积模拟更大的 batch
# 混合精度训练
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    with autocast():
        output = model(input)
        loss = criterion(output, target)

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

踩坑三:过拟合

Transformer 参数量大,容易过拟合。

表现:训练集准确率 99%,验证集只有 60%

解决

# Dropout
self.dropout = nn.Dropout(0.1)

# 权重衰减
optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.01)

# 早停
if val_loss < best_val_loss:
    best_val_loss = val_loss
    torch.save(model.state_dict(), 'best_model.pth')
    patience_counter = 0
else:
    patience_counter += 1
    if patience_counter > 5:
        break

实践中的选择

什么时候用 CNN

  • 计算资源有限
  • 文本长度较短
  • 对速度要求高
  • 不需要深层的上下文理解

什么时候用 Transformer

  • 文本长度较长
  • 上下文理解重要
  • 有充足的计算资源
  • 需要预训练模型

什么时候都不用

  • 简单规则能解决的问题
  • 数据量太少
  • 对精度要求不高

从 LeNet 到 Transformer

神经网络架构的演进,其实是对"什么重要"的理解在不断变化:

timeline title 神经网络架构演进 section 1990s LeNet : 卷积核、池化层<br/>手工设计特征 section 2010s AlexNet : 深度网络<br/>GPU 加速 VGG : 小卷积核<br/>规整结构 ResNet : 残差连接<br/>解决梯度消失 section 2020s Transformer : 注意力机制<br/>并行计算 ViT : 视觉 Transformer<br/>统一视觉语言

写在最后

模型架构这东西,没有最好的,只有最适合的。

选型之前先问自己几个问题:

  • 数据量有多大?
  • 计算资源够不够?
  • 对精度和速度的要求是什么?
  • 团队对这个架构的熟悉程度如何?

回答清楚这些问题,选择就简单了。


这次迁移花了不少时间,中间差点放弃回 CNN。但最终效果提升明显,折腾是值得的。

版权声明: 本文首发于 指尖魔法屋-从LeNet走到Transformer:神经网络架构设计笔记https://blog.thinkmoon.cn/post/19-neural-network-architecture-history/) 转载或引用必须申明原指尖魔法屋来源及源地址!