微调 Florence-2 - 微软最先进的视觉语言模型

本代码库支持我们于2024年6月24日在 Huggingface.co 上发布的一篇博客！

Florence 2 是微软于2024年6月发布的基础视觉语言模型。这个模型因其小巧的规模（0.2B 和 0.7B）以及在各种计算机视觉和视觉语言任务上的强大性能而备受青睐。

Florence 开箱即用地支持图像描述、目标检测、OCR 等功能。但是，你的任务可能不被支持，或者你可能需要控制模型在你的任务中的输出。这时你就需要对模型进行微调。

在这个代码库中，我们提供了在 DocVQA 或 The Cauldron 上微调 Florence 的代码。

安装

我们使用 UV 来管理包。UV 是一个用 Rust 编写的极快的 Python 包安装器和解析器。你可以在这里获取它。

要开始使用，请运行以下命令：

uv venv
source .venv/bin/activate
uv pip install -r requirements.txt

如果你遇到 flash-attn 的问题，可以用以下命令修复：

uv pip install -U flash-attn --no-build-isolation

获取数据

对于这个实验，我们使用 DocVQA 数据集。我们在 Hugging Face 的团队已经将其上传到 hub，所以你可以直接使用。

数据集经过预处理，便于处理。使用 HF 的 dataset 库加载后，它看起来是这样的：

from datasets import load_dataset

data = load_dataset('HuggingFaceM4/DocumentVQA')

print(data)

输出：

DatasetDict({
    train: Dataset({
        features: ['questionId', 'question', 'question_types', 'image', 'docId', 'ucsf_document_id', 'ucsf_document_page_no', 'answers'],
        num_rows: 39463
    })
    validation: Dataset({
        features: ['questionId', 'question', 'question_types', 'image', 'docId', 'ucsf_document_id', 'ucsf_document_page_no', 'answers'],
        num_rows: 5349
    })
    test: Dataset({
        features: ['questionId', 'question', 'question_types', 'image', 'docId', 'ucsf_document_id', 'ucsf_document_page_no', 'answers'],
        num_rows: 5188
    })
})

更新 Florence-2 以进行微调

为了微调 Florence-2，我们对 Florence2Seq2SeqLMOutput 类进行了一些修改。我们已经提交了这些更改的拉取请求（PR），你可以在代码中找到它们的链接。要使用修订后的版本，你需要在加载模型时指定适当的修订版本：

model = AutoModelForCausalLM.from_pretrained(
        "andito/Florence-2-large-ft", trust_remote_code=True
    ).to(device)
alternative_model = AutoModelForCausalLM.from_pretrained(
        "microsoft/Florence-2-large-ft", trust_remote_code=True, revision="refs/pr/10"
    ).to(device)

我们已经为微软的所有模型提交了必要的修复 PR。如果你想使用不同的模型，请参考我们创建的相应修订版本。

单 GPU 训练

要仅使用一个 GPU 进行训练，你可以简单地运行：

python train.py

它会自动在 DocVQA 数据集上进行训练。不建议仅使用一个 GPU 在 cauldron 上进行训练。

分布式训练

distributed_train.py 脚本允许你使用分布式数据并行来训练 Florence-2 模型，这可以在使用多个 GPU 时显著加快训练过程。以下是使用此脚本的步骤：

python distributed_train.py --dataset <dataset_name> --epochs <num_epochs> --eval-steps <evaluation_steps>

示例：

python distributed_train.py --dataset docvqa --epochs 10 --eval-steps 1000

• dataset_name：要使用的数据集名称（docvqa 或 cauldron）。
• num_epochs：要训练的轮数（默认为10）。
• evaluation_steps：训练期间评估的频率（默认为每10000步）。