Transformers 文档

图像字幕

Transformers

加入 Hugging Face 社区

并获取增强型文档体验

协作模型、数据集和空间

使用加速推理实现更快的示例

在文档主题之间切换

开始

图像字幕

图像字幕是指为给定图像预测字幕的任务。常见的现实世界应用包括帮助视障人士，帮助他们应对各种情况。因此，图像字幕通过为视障人士描述图像，帮助提高内容可访问性。

本指南将向您展示如何

微调图像字幕模型。
使用微调后的模型进行推理。

在开始之前，请确保已安装所有必要的库

pip install transformers datasets evaluate -q
pip install jiwer -q

我们鼓励您登录 Hugging Face 帐户，以便您可以将您的模型上传并与社区共享。当系统提示您时，请输入您的令牌以登录

from huggingface_hub import notebook_login

notebook_login()

加载口袋妖怪 BLIP 字幕数据集

使用 🤗 Dataset 库加载一个包含 {image-caption} 对的数据集。要使用 PyTorch 创建自己的图像字幕数据集，您可以按照此笔记本。

from datasets import load_dataset

ds = load_dataset("lambdalabs/pokemon-blip-captions")
ds

DatasetDict({
    train: Dataset({
        features: ['image', 'text'],
        num_rows: 833
    })
})

该数据集有两个特征：image 和 text。

许多图像字幕数据集包含每个图像的多个字幕。在这种情况下，一种常见的策略是在训练期间随机抽取可用字幕中的一个。

使用 train_test_split 方法将数据集的训练分割拆分为训练集和测试集

ds = ds["train"].train_test_split(test_size=0.1)
train_ds = ds["train"]
test_ds = ds["test"]

让我们可视化训练集中的一些样本。

from textwrap import wrap
import matplotlib.pyplot as plt
import numpy as np


def plot_images(images, captions):
    plt.figure(figsize=(20, 20))
    for i in range(len(images)):
        ax = plt.subplot(1, len(images), i + 1)
        caption = captions[i]
        caption = "\n".join(wrap(caption, 12))
        plt.title(caption)
        plt.imshow(images[i])
        plt.axis("off")


sample_images_to_visualize = [np.array(train_ds[i]["image"]) for i in range(5)]
sample_captions = [train_ds[i]["text"] for i in range(5)]
plot_images(sample_images_to_visualize, sample_captions)

预处理数据集

由于该数据集具有两种模态（图像和文本），因此预处理管道将预处理图像和字幕。

为此，加载与您将要微调的模型关联的处理器类。

from transformers import AutoProcessor

checkpoint = "microsoft/git-base"
processor = AutoProcessor.from_pretrained(checkpoint)

处理器将在内部预处理图像（包括调整大小和像素缩放）并对字幕进行标记。

def transforms(example_batch):
    images = [x for x in example_batch["image"]]
    captions = [x for x in example_batch["text"]]
    inputs = processor(images=images, text=captions, padding="max_length")
    inputs.update({"labels": inputs["input_ids"]})
    return inputs


train_ds.set_transform(transforms)
test_ds.set_transform(transforms)

数据集准备就绪后，您现在可以设置模型进行微调。

加载基本模型

将 “microsoft/git-base” 加载到 AutoModelForCausalLM 对象中。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(checkpoint)

评估

图像字幕模型通常使用 Rouge 得分或词错误率进行评估。对于本指南，您将使用词错误率 (WER)。

我们使用 🤗 Evaluate 库来实现这一点。有关 WER 的潜在限制和其他注意事项，请参阅本指南。

from evaluate import load
import torch

wer = load("wer")


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predicted = logits.argmax(-1)
    decoded_labels = processor.batch_decode(labels, skip_special_tokens=True)
    decoded_predictions = processor.batch_decode(predicted, skip_special_tokens=True)
    wer_score = wer.compute(predictions=decoded_predictions, references=decoded_labels)
    return {"wer_score": wer_score}

训练！

现在，您已准备好开始微调模型。您将为此使用 🤗 Trainer。

首先，使用 TrainingArguments 定义训练参数。

from transformers import TrainingArguments, Trainer

model_name = checkpoint.split("/")[1]

training_args = TrainingArguments(
    output_dir=f"{model_name}-pokemon",
    learning_rate=5e-5,
    num_train_epochs=50,
    fp16=True,
    per_device_train_batch_size=32,
    per_device_eval_batch_size=32,
    gradient_accumulation_steps=2,
    save_total_limit=3,
    eval_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=50,
    logging_steps=50,
    remove_unused_columns=False,
    push_to_hub=True,
    label_names=["labels"],
    load_best_model_at_end=True,
)

然后将它们与数据集和模型一起传递给 🤗 Trainer。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    compute_metrics=compute_metrics,
)

要开始训练，只需在 Trainer 对象上调用 train()。

trainer.train()

您应该看到训练损失随着训练的进行而平稳下降。

训练完成后，使用 push_to_hub() 方法将您的模型共享到 Hub，以便每个人都可以使用您的模型

trainer.push_to_hub()

推理

从 test_ds 中取一个样本图像来测试模型。

from PIL import Image
import requests

url = "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/pokemon.png"
image = Image.open(requests.get(url, stream=True).raw)
image

为模型准备图像。

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

inputs = processor(images=image, return_tensors="pt").to(device)
pixel_values = inputs.pixel_values

调用 generate 并解码预测。

generated_ids = model.generate(pixel_values=pixel_values, max_length=50)
generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(generated_caption)

a drawing of a pink and blue pokemon

看起来微调后的模型生成了一个相当不错的字幕！

< > 在 GitHub 上更新

←计算机视觉中的知识蒸馏文档问答→