从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/) 转载或引用必须申明原指尖魔法屋来源及源地址!
评论
使用 GitHub 账号登录后即可留言,支持 Markdown。