层次注意力文本分类模型实现：批处理与优化器权重更新问题咨询

我之前在实现层次注意力网络（HAN）做文本分类时也碰到过一模一样的问题——大文档下句子编码器要反复前向传播，批处理和优化器更新的逻辑很容易乱。下面是我踩坑后总结的可行方案：

一、批处理的高效管理 #

1. 文档级批处理+句子级批量计算 #

• 别把所有句子混在一起做全局批处理，而是按文档维度组织批次：每个批次包含N个文档，每个文档内部包含数量不等的句子。
• 先把当前批次里所有文档的句子收集起来，做一次句子编码器的批量前向传播（充分利用GPU并行能力，避免单个句子重复调用浪费资源），再按文档把句子编码结果拆分出来，送入文档编码器。
• 举个PyTorch的代码示例：

# 假设batch是文档列表，每个元素是句子张量的列表
all_sentences = []
doc_sent_counts = []
for doc in batch:
    doc_sent_counts.append(len(doc))
    all_sentences.extend(doc)
# 批量处理所有句子，得到句子编码
sent_encodings = sentence_encoder(torch.stack(all_sentences))
# 按原文档拆分编码结果
doc_sent_encodings = []
start_idx = 0
for count in doc_sent_counts:
    doc_sent_encodings.append(sent_encodings[start_idx:start_idx+count])
    start_idx += count
# 处理文档编码器
doc_encodings = [document_encoder(se) for se in doc_sent_encodings]

2. 填充与掩码的配套使用 #

• 针对长度不一的文档，在句子维度做填充时，一定要配合注意力掩码，防止文档编码器把填充的无效句子纳入注意力计算。
• 句子编码器输出后，给每个文档的句子编码添加掩码张量（标记真实句子/填充项），文档编码器的注意力层要接收这个掩码并忽略无效位置。

二、优化器权重更新的正确逻辑 #

1. 梯度累积平衡显存与效率 #

• 大文档单文档更新梯度效率太低，直接整批次更新又可能爆显存，这时候用梯度累积就很合适：
  • 设定累积步数K，每处理K个文档（或K个小批次）后，再执行一次optimizer.step()和optimizer.zero_grad()。
  • 注意累积过程中，每次前向+反向传播后不要清零梯度，让梯度自然累加；同时损失要除以累积步数，避免梯度爆炸。
• 代码示例：

accumulation_steps = 4
optimizer.zero_grad()
for idx, batch in enumerate(dataloader):
    # 前向传播计算当前批次损失
    loss = compute_loss(batch, sentence_encoder, document_encoder)
    # 损失归一化，防止梯度累积过大
    loss = loss / accumulation_steps
    loss.backward()
    # 每累积K步执行一次权重更新
    if (idx + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

2. 避免计算图内存泄漏 #

• 处理大文档时，句子编码器的计算图如果长期保留，会导致显存占用飙升。正确的做法是：
  • 每个批次处理完后，及时用del删除不再需要的中间张量，再调用torch.cuda.empty_cache()（GPU环境下）释放显存。
  • 不要在循环中存储过多历史计算图，比如不要把所有文档的损失都攒起来再反向传播，而是逐批次（或累积批次）处理。

三、额外的性能优化技巧 #

• 句子编码器预训练+阶段性冻结：如果数据集规模不大，可以先用预训练模型（比如BERT、Word2Vec）初始化句子编码器，初期冻结其权重，只训练文档编码器；待文档编码器收敛后，再解冻两者联合训练，能大幅减少计算量和显存压力。
• 动态截断/拆分长文档：对于超长文档，可设定最大句子数阈值截断超出部分，或者用滑动窗口拆分为多个子文档，分别编码后再融合结果，避免单次计算量过大。

内容的提问来源于stack exchange，提问作者Jadiel de Armas