解决Transformer模型无法处理批量大小大于1的问题：当未定义填充令牌时

1次阅读

共计 1910 个字符，预计需要花费 5 分钟才能阅读完成。

Transformer 模型在自然语言处理任务中表现出色，但在实际应用中，开发者经常遇到一个棘手的问题：当未定义填充令牌（padding token）时，模型无法处理批量大小（batch size）大于 1 的输入。这主要是因为 Transformer 模型要求输入序列的长度一致，而批量处理时不同样本的长度可能不同。

问题根源：Transformer 的自注意力机制需要计算所有位置之间的关系，因此输入序列必须具有相同的长度。如果没有填充令牌，模型无法自动对齐不同长度的序列。
实际影响：这一问题会导致推理效率低下，尤其是在生产环境中，无法利用 GPU 的并行计算能力，从而增加推理时间和资源消耗。

针对这一问题，开发者通常有以下几种解决方案：

动态填充：在批处理时动态计算最长序列长度，并为较短的序列填充填充令牌。
优点：实现简单，兼容大多数现有模型。
缺点：填充可能导致计算资源浪费，尤其是当序列长度差异较大时。
自定义批处理逻辑：在数据加载器中实现自定义的批处理逻辑，确保每个批次中的序列长度相近。
优点：减少填充带来的计算浪费。
缺点：实现复杂，可能需要调整数据加载流程。
模型配置调整：在模型定义时显式指定填充令牌，并确保模型能够正确处理填充位置。
优点：一劳永逸，适用于长期项目。
缺点：需要修改模型代码，可能影响现有部署。

以下是使用 PyTorch 实现动态填充的代码示例：

import torch
from torch.nn.utils.rnn import pad_sequence

# 假设 inputs 是一个包含多个序列的列表，每个序列是一个整数列表
inputs = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]

# 将序列转换为张量
input_tensors = [torch.tensor(x) for x in inputs]

# 动态填充
padded_inputs = pad_sequence(input_tensors, batch_first=True, padding_value=0)

print(padded_inputs)

关键点 ：pad_sequence 函数会自动计算最长序列长度，并为较短的序列填充指定的填充值（这里是 0）。

以下是使用 PyTorch 的 DataLoader 实现自定义批处理的示例：

from torch.utils.data import DataLoader, Dataset

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

def collate_fn(batch):
    # 按序列长度排序
    batch = sorted(batch, key=lambda x: len(x), reverse=True)
    # 动态填充
    padded_batch = pad_sequence([torch.tensor(x) for x in batch], batch_first=True, padding_value=0)
    return padded_batch

# 示例数据
data = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
dataset = CustomDataset(data)
dataloader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn)

for batch in dataloader:
    print(batch)