图像描述生成¶
图像描述生成（Image Captioning）是指为给定图像生成描述性文字的任务。它在现实世界中的应用包括帮助视障人士在不同情境中导航。因此，图像描述生成通过为人们描述图像，提高了内容的可访问性。
本指南将展示如何：
- 微调一个图像描述生成模型。
- 使用微调后的模型进行推理。
在开始之前，请确保安装了所有必要的库：
In [ ]:
pip install transformers datasets evaluate -q
pip install jiwer -q
我们鼓励您登录您的 Hugging Face 账户，以便上传并与社区分享您的模型。当提示时，输入您的令牌以登录：
In [ ]:
from huggingface_hub import notebook_login
notebook_login()
In [ ]:
from datasets import load_dataset
ds = load_dataset("lambdalabs/pokemon-blip-captions")
ds
数据集有两个特征：image 和 text。
许多图像描述生成数据集为每张图像包含多个描述。在这种情况下，常见的策略是在训练过程中从可用的描述中随机采样一个。
使用 train_test_split 方法将数据集的训练部分分为训练集和测试集：
In [ ]:
ds = ds["train"].train_test_split(test_size=0.1)
train_ds = ds["train"]
test_ds = ds["test"]
让我们可视化训练集中的几个样本。
In [ ]:
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)
In [ ]:
from transformers import AutoProcessor
checkpoint = "microsoft/git-base"
processor = AutoProcessor.from_pretrained(checkpoint)
处理器将内部预处理图像（包括调整大小和像素缩放）并对描述进行分词。
In [ ]:
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)
In [ ]:
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(checkpoint)
评估¶
图像描述生成模型通常使用 Rouge Score 或 Word Error Rate 进行评估。在本指南中，您将使用 Word Error Rate (WER)。
我们使用 🤗 Evaluate 库来进行评估。关于 WER 的潜在限制和其他注意事项，请参考这个指南。
In [ ]:
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}
In [ ]:
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。
In [ ]:
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_ds,
eval_dataset=test_ds,
compute_metrics=compute_metrics,
)
In [ ]:
trainer.train()
您应该看到训练损失随着训练的进行而平稳下降。
训练完成后，使用 push_to_hub() 方法将您的模型分享到 Hub，以便所有人都可以使用您的模型：
In [ ]:
trainer.push_to_hub()
推理¶
从 test_ds 中取一个样本图像来测试模型。
In [ ]:
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
准备图像以供模型使用。
In [ ]:
device = "cuda" if torch.cuda.is_available() else "cpu"
inputs = processor(images=image, return_tensors="pt").to(device)
pixel_values = inputs.pixel_values
调用 generate 并解码预测结果。
In [ ]:
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)
从结果看来，微调后的模型生成了一个相当不错的描述！