如何解决TrOCR微调模型中重复添加起始token的问题

问题背景 #

使用模型：TrOCR模型
出现场景：

• 基于NielsRogge的TrOCR微调官方示例脚本
• 自定义修改脚本（代码如下）

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")

class Dataset(Dataset):
    def __init__(self, root_dir, df, processor, max_target_length=128):
        self.root_dir = root_dir
        self.df = df
        self.processor = processor
        self.max_target_length = max_target_length

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

    def __getitem__(self, idx):
        # 获取文件名与文本
        file_name = self.df['file_name'][idx]
        text = self.df['text'][idx]
        # 预处理图像（缩放+归一化）
        image = Image.open(self.root_dir + file_name).convert("RGB")
        pixel_values = self.processor(image, return_tensors="pt").pixel_values
        # 编码文本生成标签（input_ids）
        labels = self.processor.tokenizer(text, 
                                          padding="max_length",
                                                         max_length=self.max_target_length).input_ids
        # 关键：确保PAD token被损失函数忽略
        labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
        # 返回编码结果
        return {"pixel_values": pixel_values.squeeze(), "labels": torch.tensor(labels)}

model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.vocab_size = model.config.decoder.vocab_size

model.config.eos_token_id = processor.tokenizer.sep_token_id

# python3 train.py path/to/labels  path/to/images/

环境信息：

• 系统：Linux Ubuntu（GCC 9.4.0）
• PyTorch版本：0.8.2+cu110（支持GPU）
• transformers版本：4.22.2
• Python版本：3.8.10

问题详情 #

训练和评估阶段，模型会重复添加起始token <s>（对应ID为[0,0,......,2,1,1,1]）：

1. 训练阶段compute_metrics示例：
   • 预测结果：[[0,0,506,4422,8046,2,1,1,1,1,1]]
   • 参考标签：[[0,597,2747 ...,1,1,1]]
2. 测试阶段同样存在该问题

预期结果 #

训练阶段预测结果仅含单个起始token，如[[0,506,4422,8046,2,1,1,1,1,1]]；测试阶段生成文本无重复起始token，示例：
tensor([[0,11867,405,22379,1277,..........,368,2]])
对应文本：<s>关于这个，你想拍多少就拍多少，我的看法是……</s>

解决方案 #

1. 修正标签编码逻辑 #

当前代码中，tokenizer编码文本时会自动添加起始token <s>，若原始文本或编码逻辑重复添加，就会出现双起始token。调整编码步骤，手动控制特殊token的添加：

# 编码时不自动加特殊token，手动控制
tokenized_text = self.processor.tokenizer(text, add_special_tokens=False, padding="max_length", max_length=self.max_target_length-2).input_ids
# 手动添加单个起始、结束token，再补全长度
labels = [self.processor.tokenizer.cls_token_id] + tokenized_text + [self.processor.tokenizer.sep_token_id]
# 补全到max_length
while len(labels) < self.max_target_length:
    labels.append(self.processor.tokenizer.pad_token_id)
# 处理pad token，让损失函数忽略
labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]

2. 调整模型起始token配置 #

TrOCR的decoder基于BERT，正确的起始token应为bos_token_id（部分模型中bos_token与cls_token是同一值，但显式设置更稳妥）：

# 替换原配置行
model.config.decoder_start_token_id = processor.tokenizer.bos_token_id or processor.tokenizer.cls_token_id

3. 生成阶段参数控制 #

测试或评估生成时，显式指定单个起始token作为decoder输入：

generated_ids = model.generate(pixel_values, decoder_start_token_id=processor.tokenizer.bos_token_id or processor.tokenizer.cls_token_id)

4. 升级依赖版本 #

你使用的transformers 4.22.2属于较旧版本，存在VisionEncoderDecoderModel生成逻辑的已知bug，建议升级到4.30.0+版本，同时搭配PyTorch 1.12+版本以保证兼容性。

内容的提问来源于stack exchange，提问作者Mohammed