微调 T5Gemma-2

T5Gemma 2 是来自 Google 的一系列轻量、开放权重的编码器—解码器模型，基于 Gemma 3 构建，支持多语言和多模态输入。

它在 140+ 种语言上提供最长 128K 的上下文窗口，并采用参数高效的设计选择（如权重共享的嵌入与合并注意力），因此非常适合文本生成与图像理解任务，同时体量足够小，可在笔记本电脑上运行。fin

在本教程中，我们将学习如何在 LaTeX OCR 数据集上微调编码器—解码器模型。目标是在尽量少的训练样本下获得较强表现。

1. 在 RunPod 上搭建 A100 GPU 环境

虽然可以在 Kaggle 或 Google Colab 上微调该模型，但这往往会导致会话不稳定、资源断连以及显著更慢的训练。为避免这些摩擦点并保持简单可靠的设置，我们将使用 NVIDIA A100 GPU。

您可以在 RunPod 上以每小时约 $1.39 的价格租用一块 A100，而本教程中的完整训练流程应在30 分钟以内完成。此设置可在无需与内存限制“搏斗”的情况下，提供稳定一致的性能。

首先，前往RunPod，使用最新的 PyTorch 镜像创建一个新 pod。选择1× A100 机器。

接着，编辑 pod 配置并添加名为 HF_TOKEN 的环境变量。该令牌用于：

• 从 Hugging Face 加载受限模型
• 将您微调后的模型推送回 Hugging Face Hub

当 pod 就绪后，启动笔记本并安装所需的 Python 包。请确保您使用的是最新版的 transformers。

!pip -q install -U accelerate datasets pillow sentencepiece safetensors
!pip install --quiet "transformers==5.0.0rc1"
!pip install --quiet --no-deps trl技术支持

现在，导入我们在整个笔记本中将用到的库与工具。

import torch
from datasets import load_dataset
from transformers import AutoProcessor, AutoModelForSeq2SeqLM
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer, set_seed技术支持

在训练前，我们应用一些针对 A100 GPU 优化的配置。设置随机种子可确保可复现性，启用 TF32 能在不影响 bf16 训练稳定性的情况下提升性能。

set_seed(42)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

print("CUDA:", torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
print("bf16 supported:", torch.cuda.is_available() and torch.cuda.is_bf16_supported())技术支持

CUDA: True NVIDIA A100 80GB PCIe
bf16 supported: True技术支持

2. 为小样本训练预处理 LaTeX OCR 数据集

在本教程中，我们使用 Hugging Face 上的LaTeX OCR 数据集。数据集中每个样本包含：

• 一张包含数学表达式的图像，以及
• 相应的LaTeX 源文本作为目标输出。

该数据集提供多种配置。为保持训练高效并契合我们小样本学习的目标，我们将显式控制数据集规模。

首先加载训练集与验证集。为模拟小数据微调场景，我们随机打乱数据集并选取一个小子集：

• 1,000 个样本用于训练
• 200 个样本用于验证

这样既能保证训练速度，又能让模型具备一定泛化能力。

DATASET_NAME = "full"
raw_train = load_dataset("linxy/LaTeX_OCR", name=DATASET_NAME, split="train")
raw_val   = load_dataset("linxy/LaTeX_OCR", name=DATASET_NAME, split="validation")

train_ds = raw_train.shuffle(seed=42).select(range(1000))
val_ds   = raw_val.shuffle(seed=42).select(range(200)) 

print(train_ds, val_ds)
print("Columns:", train_ds.column_names)技术支持

输出：

Dataset({
    features: ['image', 'text'],
    num_rows: 1000
}) Dataset({
    features: ['image', 'text'],
    num_rows: 200
})
Columns: ['image', 'text']技术支持

为更好地理解数据，我们来看一个训练样本。首先查看图像。

train_ds[10]["image"]技术支持

接着，查看与该图像关联的原始 LaTeX 字符串。

train_ds[10]["text"]技术支持

'G ( \\beta , \\tilde { \\mu } ) = \\left( \\frac { \\pi \\mu \\Gamma ( \\frac { \\lambda } { \\lambda + 1 } ) } { 2 \\Gamma ( \\frac { 1 } { \\lambda + 1 } ) } \\right) ^ { \\frac { 1 } { 2 \\lambda } } g _ { 0 } ( \\beta ) g _ { S } ( \\beta , Z ) ,'技术支持

最后，我们可以在笔记本中直接渲染该 LaTeX 表达式，验证文本与图像内容是否匹配。

from IPython.display import display, Math, Latex

latex = train_ds[10]["text"]
display(Math(latex))技术支持

3. 初始化 T5Gemma-2 多模态编码器—解码器

数据集准备就绪后，下一步是加载T5Gemma-2 模型及其对应的处理器。我们将使用270M–270M 变体，它在能力与效率之间取得良好平衡，并可轻松适配消费级 GPU。

MODEL_ID = "google/t5gemma-2-270m-270m"

processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True)
model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,   # A100 -> bf16
    device_map="auto",
)技术支持

分词器通过处理器访问。某些编码器—解码器模型默认未定义填充符，因此如有必要我们会添加一个，并相应地调整模型的词嵌入大小。

tokenizer = processor.tokenizer
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})
    model.resize_token_embeddings(len(tokenizer))技术支持

针对 LaTeX OCR，我们使用简短的文本提示与输入图像配对来引导模型。该提示清晰地指定任务并约束输出格式。

最后，我们为编码器输入与解码器目标定义最大序列长度。这些取值足以覆盖大多数数学表达式，同时控制内存占用。

PROMPT = "<start_of_image> Convert this image to LaTeX. Output only LaTeX."
MAX_INPUT_LEN  = 128
MAX_TARGET_LEN = 256技术支持

4. 运行基线推理：零样本表现检查

在开始微调之前，先用预训练的 T5Gemma-2 模型运行一次基线推理，有助于我们了解模型在 LaTeX OCR 任务上未经过任务特定训练时的表现，并为后续改进提供参考。

流程简述如下：

1. 从训练数据集中选择一个样本图像。
2. 将图像与文本提示组合为所需的张量格式，准备模型输入。
3. 在评估模式下运行推理，并禁用梯度计算以节省内存并加快生成速度。
4. 使用有限束数的束搜索与适度的重复惩罚，以获得更干净的输出。
5. 生成完成后，将生成的标记解码回文本。

image = train_ds[20]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)

print("\n--- Model output ---")
print(pred)技术支持

预训练模型的示例输出：

--- Model output ---
Use this command to change this image's format:\usepackage[T1]{fontenc}\usepackage[utf8]{inputenc}\DeclareUnicodeUTF8{1234567890123}\begin{document}<tex>$\begin{equation}\begin{aligned}P=&\frac{1+(-)\frac{1}{2}\left(F_{1}(-)\frac{F_{2}(-技术支持

现在，来看同一图像的真实（ground-truth）LaTeX。

train_ds[20]["text"]技术支持

真实标注：

'{ \\cal P } = \\frac { 1 + ( - ) ^ { F } { \\cal I } _ { 4 } ( - ) ^ { F _ { L } } } { 2 } .'技术支持

如预期，预训练模型并未生成正确的 LaTeX 转写，而是输出通用的 LaTeX 模板与不相关片段。这是正常现象，因为基础模型未针对 LaTeX OCR 进行特定训练。

5. 构建自定义图文数据整理器（Data Collator）

本节我们将构建一个自定义的图文数据整理器，以便在微调时正确地批处理图像、提示与 LaTeX 目标。由于 T5Gemma-2 是多模态编码器—解码器模型，整理器对确保图像与文本正确对齐并以预期格式传入模型至关重要。

具体而言，该整理器会：

• 从不同可能格式加载图像并统一转换为 RGB 表示
• 确保每个训练样本恰好包含一张图像，并封装为处理器期望的结构
• 为每张图像附加固定指令提示，清晰定义 OCR 任务
• 对 LaTeX 目标单独分词，并进行受控截断与填充
• 在标签中屏蔽填充符，使其不计入训练损失
• 避免对输入端进行截断，以防止图像—标记不匹配

上述步骤共同确保在 LaTeX OCR 数据上微调时的稳定训练、正确的损失计算以及可靠的收敛。

from typing import Any, Dict, List
import torch
from PIL import Image as PILImage

tokenizer = processor.tokenizer
tokenizer.padding_side = "right"

if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})

pad_id = tokenizer.pad_token_id

PROMPT = "<start_of_image> Convert this equation image to LaTeX. Output only LaTeX."
MAX_TARGET_LEN = 256

def collate_fn(examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
    images, prompts, targets = [], [], []

    for ex in examples:
        im = ex["image"]
        if isinstance(im, PILImage.Image):
            im = im.convert("RGB")
        elif isinstance(im, dict) and "path" in im:
            im = PILImage.open(im["path"]).convert("RGB")
        else:
            raise ValueError(f"Unexpected image type: {type(im)}")

        # IMPORTANT: one image per sample -> nested list
        images.append([im])

        prompts.append(PROMPT)
        targets.append(ex["text"])

    # ✅ NO truncation here (prevents image-token mismatch)
    model_inputs = processor(
        text=prompts,
        images=images,
        padding=True,
        truncation=False,
        return_tensors="pt",
    )

    labels = tokenizer(
        targets,
        padding=True,
        truncation=True,
        max_length=MAX_TARGET_LEN,
        return_tensors="pt",
    )["input_ids"]

    labels[labels == pad_id] = -100
    model_inputs["labels"] = labels
    return model_inputs技术支持

6. 训练配置与优化设置

本节我们定义在 LaTeX OCR 任务上微调 T5Gemma-2 的训练配置与优化设置。该设置有意保持简单，并针对单块 A100、使用小型数据集与单轮训练进行加速优化。

为减少开销并提速，我们禁用评估与检查点保存，仅专注于高效的前向与反向传播。

from transformers import Seq2SeqTrainingArguments

args = Seq2SeqTrainingArguments(
    output_dir="t5gemma2-latex-ocr-1k",

    # --- core training ---
    num_train_epochs=1,
    per_device_train_batch_size=8, 
    gradient_accumulation_steps=1,

    learning_rate=1e-4,
    warmup_steps=15,
    lr_scheduler_type="linear",

    # --- precision / speed ---
    bf16=True,    
    fp16=False,
    tf32=True,     

    # --- memory ---
    gradient_checkpointing=True,

    # --- stop extra work (this is the big speed win) ---
    eval_strategy="no",           
    predict_with_generate=False,          
    save_strategy="no",                 
    report_to="none",

    # --- dataloader ---
    dataloader_num_workers=0,         
    remove_unused_columns=False,

    # --- logging ---
    logging_steps=10,
)技术支持

最后，初始化 Seq2SeqTrainer。我们显式传入自定义数据整理器，以便训练器在训练期间正确构建包含图像、指令提示与 LaTeX 目标序列的多模态批次。

trainer = Seq2SeqTrainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=collate_fn,
)技术支持

7. 在 LaTeX OCR 上微调 T5Gemma-2

一切就绪后，我们就可以开始微调模型了。通过对训练器的一次调用即可启动训练。

trainer.train()技术支持

在训练过程中，损失会逐步下降，表明模型正在学习更好地将方程图像映射到相应的 LaTeX 表达。

即便是小数据集且仅训练一轮，模型也能迅速适应 OCR 任务。

8. 微调后的评估

微调完成后，我们分别在一个训练样本和一个验证样本上重新运行推理，观察模型输出的变化。

与基线相比（主要生成通用 LaTeX 模板），微调后的模型能够生成结构化的 LaTeX，并与目标方程的形状与符号更加接近。

我们先在之前相同的训练样本上测试。图像经处理器处理，模型生成标记，再解码回 LaTeX。

# pick a sample
image = train_ds[10]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)
print(pred)技术支持

如您所见，我们现在得到的是正确的 LaTeX 输出，而非随机或无关文本。在该训练样本中，预测与真实值在大体上对齐，剩余差错多发生在分数、下标与少量符号位置上。

G ( \beta , \tilde { \mu } ) = \left( \frac { \pi \mu \Gamma } { 2 \Gamma } \lambda _ { + 1 } ^ { \lambda } \right) \right) ^ { \frac { 1 } { 9 \Gamma } g _ { 0 } ( \beta ) g _ { S } ( \theta , Z ) , \right) , \qquad G ( \bar { \技术支持

接下来我们在一个验证样本上测试以检查泛化能力。

image = val_ds[10]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)
print(pred)技术支持

在该验证样本中，模型仍然遵循正确的 LaTeX 结构与符号，尽管有时会在括号位置、括号内项以及较长表达式上出现错误。

f ( p , p ^ { \prime } ) = \ln \left\{ \frac { ( p _ { i } - p _ { j } ) ^ { 2 } } { \left( \frac { \psi ( p _ j ) - g ( p ; p _ { s } ) \psi ( \psi _ { i ] } ) \right\} . . . g ( \psi ; \psi ) - g \left( p ; \psi _技术支持

与真实值对比，整体结构显然是对齐的，模型输出的是接近的近似而非不相干的内容。

print(val_ds[10]["text"])技术支持

f ( p , p ^ { \prime } ) = \ln \left\{ \frac { ( p _ { i } - p _ { j } ) ^ { 2 } } { ( p _ { i } + p _ { j } ) ^ { 2 } } \right\} \left[ \psi ( p _ { j } ) - g ( p _ { i } , p _ { j } ) \psi ( p _ { i } ) \right] .技术支持

总体而言，微调后结果有明显提升。模型不再套用通用 LaTeX 模板，而是生成与数据集目标高度相似的“类方程”LaTeX，即使只用少量训练集与短时间微调亦如此。

9. 保存与发布微调后的 T5Gemma-2 模型

训练完成后，第一步是将微调模型本地保存，以便后续复用进行推理。

trainer.save_model()技术支持

接着，我们将模型推送到 Hugging Face Hub，以便他人访问、复用并在此基础上继续构建。

trainer.push_to_hub()技术支持

在检查仓库文件时，您可能会注意到，通过训练器推送模型时，处理器配置并不总会被包含在内。

由于处理器对同时正确处理图像与文本是必需的，我们会将其单独显式推送，确保后续加载和使用模型无需额外配置。

processor.push_to_hub(repo_id="kingabzpro/t5gemma2-latex-ocr-1k")技术支持

完成这一步后，仓库中已包含加载模型与处理器所需的一切，且只需一次调用即可。您现在可以访问 Hugging Face 上的kingabzpro/t5gemma2-latex-ocr-1k，获取微调后的模型并开始用于 LaTeX OCR 或进一步试验。

10. 使用 Pipelines 加载模型进行推理

既然微调后的模型已发布到 Hugging Face Hub，我们就可以直接使用 pipeline API 加载进行推理。这是最简单的测试方式，无需手动处理处理器、分词器或生成逻辑。

我们从 Hub 加载模型并创建一个 image-text-to-text 管道：

from transformers import pipeline

generator = pipeline(
    "image-text-to-text",
    model="kingabzpro/t5gemma2-latex-ocr-1k",
)技术支持

接下来，我们使用与之前相同的指令提示，在一个验证样本上运行推理。

generator(
    val_ds[10]["image"],
    text="<start_of_image> Convert this image to LaTeX. Output only LaTeX.",
    generate_kwargs={"do_sample": False, "max_new_tokens": 100},
)技术支持

可以看到，输出已经非常接近正确的 LaTeX。

[{'input_text': '<start_of_image> Convert this image to LaTeX. Output only LaTeX.',
  'generated_text': '<start_of_image> Convert this image to LaTeX. Output only LaTeX.f ( p , p ^ { \\prime } ) = \\ln \\left\\{ \\begin{array} { \\begin{array} { \\begin{array} { \\begin{array} { \\begin{array} { \\end{array} \\right\\} \\begin{array} { \\begin{array} { \\begin{array} { \\end{array} \\right\\} \\begin{array} { \\begin{array} { \\begin{array} {'}]技术支持

我们再试一个验证样本，并对输出进行后处理，仅保留 LaTeX 字符串。

preds = generator(
    val_ds[30]["image"],
    text="<start_of_image> Convert this image to LaTeX. Output only LaTeX.",
    generate_kwargs={"do_sample": False, "max_new_tokens": 100},
)

prompt = preds[0]["input_text"]
gen = preds[0]["generated_text"]

# remove the prompt if the model echoed it
if gen.startswith(prompt):
    gen = gen[len(prompt):]

# remove any leftover special tokens / separators
gen = gen.replace("<start_of_image>", "").strip()
if gen.startswith("."):
    gen = gen[1:].strip()

print("\nCLEAN PREDICTED LaTeX:\n", gen)技术支持

这次的结果干净，可直接作为 LaTeX 使用：

CLEAN PREDICTED LaTeX:
 T _ { M N } = \left\{ g N \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ {技术支持

至此，模型即可作为可插即用的 LaTeX OCR 系统。您可以将其部署在 API 背后、集成到文档处理流程，或继续用更多数据进行微调以获得更高准确率。

如果在运行上述代码时遇到问题，请查看辅助笔记本。其中包含完整代码及各步骤输出，可引导您完成整个流程。

结语

刚开始训练该模型时，我将其当作带图像编码器的大型语言模型来处理。多次失败后我意识到，这种方法并不适用于序列到序列的编码器—解码器模型。

我不得不重新思考整个设置，包括数据整理器、训练器与训练参数，甚至推理的执行方式。

在本教程中，我们走完了在 LaTeX OCR 任务上微调 T5Gemma-2 的端到端流程：从环境搭建与数据集检视，到自定义数据整理、高效训练与训练后评估。

使用小型数据集与单块 A100 GPU，我们展示了多模态编码器—解码器模型如何快速学会从方程图像生成结构化且有意义的 LaTeX。

最终，微调后的模型显著超越了通用模板输出，能够生成与真实值高度匹配的“类方程”LaTeX，展示出现代开放模型在真实世界 OCR 与文档理解任务中的微调可达性与有效性。

如果您希望查看更多 LLM 微调的动手示例，建议查看《使用 Llama 3 进行微调》课程。