一个图像字幕生成代码库

这是一个用于图像字幕生成研究的代码库。

它支持：

• 自我批判训练，来自 Self-critical Sequence Training for Image Captioning
• 自下而上的特征提取，来自 ref。
• 测试时集成
• 多GPU训练。（在PyTorch Lightning的帮助下，现在支持DistributedDataParallel，详情请参阅ADVANCED.md）
• Transformer字幕生成模型。

一个简单的演示Colab笔记本可以在这里找到。

需求

• Python 3
• PyTorch 1.3+（以及torchvision）（测试版本为1.13）
• cider（已作为子模块添加）
• coco-caption（已作为子模块添加）（请务必按照coco-caption/README.md中的初始化步骤操作）
• yacs
• lmdbdict
• 可选：pytorch-lightning（测试版本为2.0）

安装

如果您在运行tools中的训练脚本时遇到困难，可以尝试将此仓库安装为Python包：

python -m pip install -e .

预训练模型

请查看MODEL_ZOO.md。

如果您只想进行评估，可以在下载预训练模型（以及预训练的ResNet101或预计算的自下而上特征，详见data/README.md）后，按照此处的说明进行操作。

在COCO/Flickr30k数据集上训练您自己的网络

准备数据。

我们现在同时支持Flickr30k和COCO数据集。详细信息请参阅data/README.md。（注意：后续部分假设使用COCO数据集；使用Flickr30k应该也很简单。）

开始训练

$ python tools/train.py --id fc --caption_model newfc --input_json data/cocotalk.json --input_fc_dir data/cocotalk_fc --input_att_dir data/cocotalk_att --input_label_h5 data/cocotalk_label.h5 --batch_size 10 --learning_rate 5e-4 --learning_rate_decay_start 0 --scheduled_sampling_start 0 --checkpoint_path log_fc --save_checkpoint_every 6000 --val_images_use 5000 --max_epochs 30

或者

$ python tools/train.py --cfg configs/fc.yml --id fc

训练脚本会将检查点保存到--checkpoint_path指定的文件夹中（默认为log_$id/）。默认情况下，为了节省磁盘空间，只保存验证集上表现最好的检查点和最新的检查点。您也可以设置--save_history_ckpt为1，以保存每一个检查点。

要恢复训练，您可以指定--start_from选项为保存infos.pkl和model.pth的路径（通常您可以将--start_from和--checkpoint_path设置为同一路径）。

要查看训练曲线或验证曲线，可以使用TensorBoard。损失历史会自动保存到--checkpoint_path中。

当前命令使用了计划采样，您也可以将--scheduled_sampling_start设置为-1来关闭计划采样。

如果您希望在训练过程中除了验证交叉熵损失外，还评估BLEU/METEOR/CIDEr分数，请使用--language_eval 1选项，但别忘了拉取子模块coco-caption。

对于所有参数，您可以将其指定在一个yaml文件中，并使用--cfg来加载该yaml文件中的配置。如果命令行中的配置与yaml文件中的配置冲突，命令行中的配置将优先。

更多选项请参阅opts.py。

使用自我批判训练

首先，您需要对数据集进行预处理，并生成用于计算CIDEr分数的缓存：

$ python scripts/prepro_ngrams.py --input_json data/dataset_coco.json --dict_json data/cocotalk.json --output_pkl data/coco-train --split train

然后，从使用交叉熵训练的预训练模型中复制模型。（这不是必须的，只是为了备份）

$ bash scripts/copy_model.sh fc fc_rl

接着：

$ python tools/train.py --id fc_rl --caption_model newfc --input_json data/cocotalk.json --input_fc_dir data/cocotalk_fc --input_att_dir data/cocotalk_att --input_label_h5 data/cocotalk_label.h5 --batch_size 10 --learning_rate 5e-5 --start_from log_fc_rl --checkpoint_path log_fc_rl --save_checkpoint_every 6000 --language_eval 1 --val_images_use 5000 --self_critical_after 30 --cached_tokens coco-train-idxs --max_epoch 50 --train_sample_n 5

或者

$ python tools/train.py --cfg configs/fc_rl.yml --id fc_rl

您会看到CIDEr分数会有显著提升，:)。

关于训练的一些注意事项。 从第30个epoch开始进行自我批判训练，经过60万次迭代后（包括之前的30个epoch），CIDEr分数可以达到1.05。

生成图像字幕

在原始图像上评估

注意：这不适用于使用自下而上特征训练的模型。现在将您感兴趣的所有图像放入一个文件夹中，例如blah，然后运行评估脚本：

$ python tools/eval.py --model model.pth --infos_path infos.pkl --image_folder blah --num_images 10

这会告诉评估脚本从给定文件夹中最多处理10张图像。如果您有大显存的GPU，可以通过增加batch_size来加快评估速度。使用--num_images -1来处理所有图像。评估脚本会在vis文件夹中生成一个vis.json文件，您可以使用提供的HTML界面来可视化结果：

$ cd vis
$ python -m SimpleHTTPServer

现在在浏览器中访问localhost:8000，您应该可以看到预测的字幕。

在Karpathy的测试集上评估

$ python tools/eval.py --dump_images 0 --num_images 5000 --model model.pth --infos_path infos.pkl --language_eval 1 

默认的评估分割是测试集。默认的推理方法是贪婪解码（--sample_method greedy），如果想从后验分布中采样，可以设置--sample_method sample。

束搜索。束搜索可以将贪婪解码序列的性能提高约5%。不过，这种方法的成本稍高。要启用束搜索，可以使用--beam_size N，N应大于1。

在COCO测试集上评估

$ python tools/eval.py --input_json cocotest.json --input_fc_dir data/cocotest_bu_fc --input_att_dir data/cocotest_bu_att --input_label_h5 none --num_images -1 --model model.pth --infos_path infos.pkl --language_eval 0

您可以从链接下载预处理好的文件cocotest.json、cocotest_bu_att和cocotest_bu_fc。

杂项

使用 CPU。当前代码默认使用 GPU；甚至没有切换选项。如果有人非常需要 CPU 版本的模型，请提交一个 issue；我可以尝试创建一个 CPU 检查点，并修改 eval.py 以在 CPU 上运行模型。不过，使用 CPU 训练该模型并无意义。

在其他数据集上训练。如果你能为自己的数据集创建一个类似 dataset_coco.json 的文件，那么移植应该非常简单。

在线演示。目前暂不支持。欢迎提交 Pull Request。

更高级的功能：

请查看 ADVANCED.md。

参考文献

如果你觉得这个仓库对你有帮助，请考虑引用（完全自愿）：

@article{luo2018discriminability,
  title={用于训练描述性标题的可区分性目标},
  author={Luo, Ruotian and Price, Brian and Cohen, Scott and Shakhnarovich, Gregory},
  journal={arXiv 预印本 arXiv:1803.04376},
  year={2018}
}

当然，也请引用你所使用的模型的原始论文（参考文献可在模型文件中找到）。

致谢

感谢原始的 neuraltalk2 以及优秀的 PyTorch 团队。

ImageCaptioning.pytorch 快速上手指南

本指南旨在帮助开发者快速部署并使用 ImageCaptioning.pytorch 进行图像描述（Image Captioning）的研究与训练。该项目基于 PyTorch，支持自临界序列训练（Self-critical）、Bottom-up 特征、Transformer 模型及多 GPU 分布式训练。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.x
• PyTorch: 1.3+ (推荐使用 1.13 及以上版本，需同时安装 torchvision)
• 可选依赖: pytorch-lightning (用于多 GPU 分布式训练，测试版本 2.0)

前置依赖说明： 项目已包含 cider 和 coco-caption 作为子模块。克隆仓库后，务必初始化子模块并遵循 coco-caption/README.md 中的初始化步骤（通常需要下载 Java 和相关脚本），否则无法计算 CIDEr 等评估指标。

git clone --recursive https://github.com/ruotianluo/ImageCaptioning.pytorch.git
cd ImageCaptioning.pytorch
# 按照 coco-caption/README.md 完成后续初始化

其他 Python 依赖可通过以下方式安装（建议使用国内镜像源加速）：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
# 或者手动安装核心依赖
pip install yacs lmdbdict -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

如果您在运行 tools 目录下的训练脚本时遇到困难，建议将本项目作为 Python 包进行安装：

python -m pip install -e .

基本使用

1. 数据准备

本项目支持 COCO 和 Flickr30k 数据集。以下示例默认使用 COCO 数据集。 请参考 data/README.md 下载预处理好的数据文件（如 cocotalk.json, cocotalk_fc, cocotalk_att, cocotalk_label.h5 等）。

2. 开始训练

您可以直接通过命令行参数或配置文件启动训练。以下是最基础的训练命令示例，使用交叉熵损失函数训练一个 newfc 模型：

python tools/train.py --id fc --caption_model newfc --input_json data/cocotalk.json --input_fc_dir data/cocotalk_fc --input_att_dir data/cocotalk_att --input_label_h5 data/cocotalk_label.h5 --batch_size 10 --learning_rate 5e-4 --learning_rate_decay_start 0 --scheduled_sampling_start 0 --checkpoint_path log_fc --save_checkpoint_every 6000 --val_images_use 5000 --max_epochs 30

或者使用配置文件（更简洁）：

python tools/train.py --cfg configs/fc.yml --id fc

• 断点续训: 添加 --start_from log_fc 参数（路径需包含 infos.pkl 和 model.pth）即可从上次保存的状态继续训练。
• 监控: 训练日志和 Loss 曲线会自动保存在 --checkpoint_path 指定的文件夹中，可使用 TensorBoard 查看。

3. 生成图像描述 (评估)

训练完成后，您可以使用预训练模型或自己的模型对图像生成描述。

对本地文件夹中的图片进行评估： 将待测试图片放入一个文件夹（例如 blah），然后运行：

python tools/eval.py --model model.pth --infos_path infos.pkl --image_folder blah --num_images 10

• --num_images: 设置要处理的图片数量，设为 -1 处理所有图片。
• 结果将生成在 vis/vis.json，可通过启动简易 HTTP 服务器在浏览器查看可视化结果：

  cd vis
  python -m SimpleHTTPServer
  # 访问 http://localhost:8000
  

在标准测试集上评估 (含 BLEU/CIDEr 分数)：

python tools/eval.py --dump_images 0 --num_images 5000 --model model.pth --infos_path infos.pkl --language_eval 1

• 束搜索 (Beam Search): 添加 --beam_size N (N>1) 可提升生成质量，但会增加计算开销。
• 采样方法: 默认为贪婪解码 (greedy)，如需从后验分布采样，设置 --sample_method sample。

提示: 若需体验快速演示，可直接访问项目提供的 Colab Notebook。