微调 Florence-2 - 微软的尖端视觉语言模型

发布于2024年6月24日
在 GitHub 上更新

Florence-2 是微软于2024年6月发布的基础视觉-语言模型。该模型因其小尺寸(0.2B 和 0.7B)以及在各种计算机视觉和视觉-语言任务上的强大性能而备受关注。

Florence 开箱即用地支持许多任务:图像标注、目标检测、OCR 等。但是,您的任务或领域可能不受支持,或者您可能希望更好地控制模型在您的任务中的输出。这时您就需要进行微调。

在这篇文章中,我们将展示一个在 DocVQA 上微调 Florence 的例子。作者报告说 Florence 2 可以执行视觉问答 (VQA),但发布的模型不包括 VQA 功能。让我们看看我们能做些什么!

预训练细节和架构

VLM Structure
Florence-2 架构

无论执行何种计算机视觉任务,Florence-2 都将问题表述为序列到序列任务。Florence-2 接收图像和文本作为输入,并生成文本作为输出。该模型结构简单。它使用 DaViT 视觉编码器将图像转换为视觉嵌入,并使用 BERT 将文本提示转换为文本和位置嵌入。然后,生成的嵌入由标准的编码器-解码器 Transformer 架构处理,生成文本和位置标记。Florence-2 的强大之处并非源于其架构,而是源于其预训练所使用的大规模数据集。作者指出,领先的计算机视觉数据集通常包含有限的信息——WIT 仅包含图像/标题对,SA-1B 仅包含图像和相关的分割掩码。因此,他们决定构建一个新的 FLD-5B 数据集,其中包含每张图像的广泛信息——边界框、掩码、标题和接地。数据集创建过程大部分是自动化的。作者使用现成的任务特定模型和一套启发式方法和质量检查来清理获得的结果。最终得到一个新数据集,包含1.26亿张图像的50亿多条标注,该数据集用于预训练 Florence-2 模型。

VQA上的原始性能

我们尝试了各种方法来调整模型以适应 VQA(视觉问答)响应。我们发现最有效的方法是区域到描述的提示,尽管它与 VQA 任务并不完全一致。图像标注提供了图像的描述性信息,但不支持直接的问答输入。我们还测试了几种“不受支持”的提示,例如“”、“”和“”。不幸的是,这些尝试都产生了无法使用的结果。

微调后在 DocVQA 上的性能

我们使用 DocVQA 数据集的标准指标 Levenshtein 相似度来衡量性能。在微调之前,模型预测与验证集上的真实标签之间的相似度为0,因为输出与真实标签不接近。在用训练集微调七个 epoch 后,验证集上的相似度得分提高到57.0。我们创建了一个 🤗 空间来演示微调后的模型。尽管该模型在 DocVQA 方面表现良好,但在通用文档理解方面仍有改进空间。然而,它成功地完成了任务,展示了 Florence-2 在下游任务中进行微调的潜力。为了开发一个出色的 VQA 模型,我们建议使用 The Cauldron 进一步微调 Florence-2。我们已经在 我们的 GitHub 页面上提供了必要的代码。

为了提供一个具体的例子,下面我们提供了微调前后两个推理结果。您也可以在这里试用该模型。

Before After Fine-tuning
微调前后

微调细节

在预训练中,作者对基础模型使用了 2048 的批量大小,对大型模型使用了 3072 的批量大小。他们还描述了与冻结图像编码器相比,解冻图像编码器进行微调时性能的提升。

我们使用低得多的资源配置进行了实验,以探索模型在更受限制的微调环境中的能力。我们在 Colab 的单个 A100 GPU 上冻结了视觉编码器并使用了 6 的批量大小,或者在 T4 上使用了 1 的批量大小。同时,我们还进行了一项使用更多资源的实验,以 64 的批量大小微调整个模型。此训练过程在配备 8 个 H100 GPU 的集群上花费了 70 分钟。此训练好的模型可在此处找到

在每种情况下,我们都发现 1e-6 的小学习率对训练有益。如果学习率较大,模型会很快过拟合训练集。

代码演练

如果您想继续学习,可以在此处找到我们的 Colab 微调笔记本,其中包含在 DocVQA 数据集上的检查点。让我们从安装依赖项开始。

!pip install -q datasets flash_attn timm einops

从 Hugging Face Hub 加载 DocVQA 数据集。

import torch
from datasets import load_dataset 

data = load_dataset("HuggingFaceM4/DocumentVQA")

我们可以使用 transformers 库中的 AutoModelForCausalLMAutoProcessor 类加载模型和处理器。我们需要传递 trust_remote_code=True,因为该模型使用了自定义代码——它尚未原生集成到 transformers 中。我们还将冻结视觉编码器以降低微调成本。

from transformers import AutoModelForCausalLM, AutoProcessor
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 

model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Florence-2-base-ft",
    trust_remote_code=True,
    revision='refs/pr/6'
).to(device) 
processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base-ft", 
    trust_remote_code=True, revision='refs/pr/6')

for param in model.vision_tower.parameters():
  param.is_trainable = False

现在让我们来微调模型!我们将构建一个训练 PyTorch 数据集,在该数据集中,我们将向数据集中的每个问题添加一个 前缀。

import torch from torch.utils.data import Dataset 

class DocVQADataset(Dataset): 

    def __init__(self, data): 
        self.data = data
        
    def __len__(self): 
        return len(self.data)
        
    def __getitem__(self, idx):
        example = self.data[idx]
        question = "<DocVQA>" + example['question'] 
        first_answer = example['answers'][0]
        image = example['image'].convert("RGB")
        return question, first_answer, image

我们现在将构建数据整理器,它从数据集样本构建训练批次,然后开始训练。在具有 40GB 内存的 A100 上,我们可以放入 6 个示例。如果您在 T4 上训练,可以使用 1 的批量大小。

import os 
from torch.utils.data import DataLoader
from tqdm import tqdm 
from transformers import AdamW, get_scheduler

def collate_fn(batch): 
    questions, answers, images = zip(*batch)
    inputs = processor(text=list(questions), images=list(images), return_tensors="pt", padding=True).to(device)
    return inputs, answers 

train_dataset = DocVQADataset(data['train'])
val_dataset = DocVQADataset(data['validation']) 
batch_size = 6
num_workers = 0

train_loader = DataLoader(train_dataset, batch_size=batch_size, 
                          collate_fn=collate_fn, num_workers=num_workers, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, 
                          collate_fn=collate_fn, num_workers=num_workers)

我们现在可以训练模型了。

epochs = 7
optimizer = AdamW(model.parameters(), lr=1e-6)
num_training_steps = epochs * len(train_loader)

lr_scheduler = get_scheduler(name="linear", optimizer=optimizer, 
                              num_warmup_steps=0, num_training_steps=num_training_steps,)

for epoch in range(epochs): 
    model.train() 
    train_loss = 0
    i = -1
    for inputs, answers in tqdm(train_loader, desc=f"Training Epoch {epoch + 1}/{epochs}"):
        i += 1
        input_ids = inputs["input_ids"]
        pixel_values = inputs["pixel_values"] 
        labels = processor.tokenizer(text=answers, return_tensors="pt", padding=True, return_token_type_ids=False).input_ids.to(device)
        outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        train_loss += loss.item()
    avg_train_loss = train_loss / len(train_loader)
    print(f"Average Training Loss: {avg_train_loss}")

    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in tqdm(val_loader, desc=f"Validation Epoch {epoch + 1}/{epochs}"):
            inputs, answers = batch
            input_ids = inputs["input_ids"]
            pixel_values = inputs["pixel_values"]
            labels = processor.tokenizer(text=answers, return_tensors="pt", padding=True, return_token_type_ids=False).input_ids.to(device)
            outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=labels)
            loss = outputs.loss
            val_loss += loss.item()

      print(val_loss / len(val_loader))

您可以通过对两个对象调用 save_pretrained() 来保存模型和处理器。完全微调的模型在此处,演示在此处

结论

在这篇文章中,我们展示了 Florence-2 可以有效地微调到自定义数据集,在短时间内在新任务上取得了令人印象深刻的性能。这种能力对于那些希望在设备上部署这个小型模型或在生产环境中经济高效地使用它的人来说尤其有价值。我们鼓励开源社区利用本微调教程,探索 Florence-2 在各种新任务中的卓越潜力!我们迫不及待地想在 🤗 Hub 上看到您的模型!

有用资源

我们感谢 Pedro Cuenca 对本博客文章的评论。

社区

注册登录 发表评论