一个图像字幕生成代码库

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

它支持：

• 自我批判训练，来自 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万次迭代后，CIDEr分数上升到1.05（包括之前的30个epoch的预训练）。

生成图像字幕

在原始图像上评估

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

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

这会告诉评估脚本从给定文件夹中最多处理10张图片。如果您有一个大显卡，可以通过增加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 模型，请提交一个问题；我可能会创建一个 CPU 检查点，并修改 eval.py 以在 CPU 上运行模型。不过，用 CPU 训练模型并无意义。

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

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

更高级的功能：

请查看 ADVANCED.md。

参考文献

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

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

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

致谢

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

self-critical.pytorch 快速上手指南

self-critical.pytorch 是一个基于 PyTorch 的图像描述（Image Captioning）研究代码库，支持自临界序列训练（Self-critical Sequence Training）、Bottom-up 特征、Transformer 模型及多 GPU 分布式训练。

环境准备

系统要求

• 操作系统: Linux / macOS (Windows 需自行适配)
• Python: 3.x
• PyTorch: 1.3+ (推荐 1.13+)，需同时安装 torchvision
• 可选依赖: pytorch-lightning (用于分布式训练，测试版本 2.0)

前置依赖

项目已包含部分子模块，但需手动初始化：

1. cider: 已作为子模块添加。
2. coco-caption: 已作为子模块添加，必须执行初始化步骤。

   git submodule update --init --recursive
   cd coco-caption
   # 按照 coco-caption/README.md 中的步骤下载数据并初始化
   

3. 其他 Python 包:

   pip install yacs lmdbdict
   

提示: 国内用户建议使用清华或阿里镜像源加速安装： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple <package_name>

安装步骤

如果遇到运行脚本困难，建议将本项目作为 Python 包进行安装：

python -m pip install -e .

基本使用

以下流程以在 COCO 数据集上训练和评估为例。

1. 数据准备

本项目支持 COCO 和 Flickr30k 数据集。请参照 data/README.md 下载并预处理数据。 注意：后续命令默认假设已准备好 COCO 数据文件（如 cocotalk.json, cocotalk_fc 等）。

2. 开始训练 (交叉熵损失)

首先使用标准的交叉熵损失训练基础模型：

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

训练检查点将保存在 log_fc/ 目录下。可使用 TensorBoard 查看训练曲线。

3. 自临界训练 (Self-Critical Training)

在基础模型训练完成后，可进一步使用自临界策略提升 CIDEr 分数。

步骤 A: 预处理 N-grams 缓存

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

步骤 B: 复制预训练模型

bash scripts/copy_model.sh fc fc_rl

步骤 C: 启动强化学习训练

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

4. 生成与评估图像描述

评估单张或多张本地图片: 将图片放入文件夹（例如 blah），运行：

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

生成的可视化结果位于 vis/vis.json，可通过以下命令在浏览器查看：

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

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

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

注：默认使用贪婪解码 (greedy)。如需开启束搜索 (Beam Search) 以提升性能，添加参数 --beam_size N (N > 1)。