Skip to content

foamliu/Image-Captioning-PyTorch

Repository files navigation

图像中文描述

图像中文描述 + 视觉注意力的 PyTorch 实现。

Show, Attend, and Tell 是令人惊叹的工作,这里是作者的原始实现。

这个模型学会了“往哪瞅”:当模型逐词生成标题时,模型的目光在图像上移动以专注于跟下一个词最相关的部分。

依赖

  • Python 3.5
  • PyTorch 0.4

数据集

使用 AI Challenger 2017 的图像中文描述数据集,包含30万张图片,150万句中文描述。训练集:210,000 张,验证集:30,000 张,测试集 A:30,000 张,测试集 B:30,000 张。

image

下载点这里:图像中文描述数据集,放在 data 目录下。

网络结构

image

用法

数据预处理

提取210,000 张训练图片和30,000 张验证图片:

$ python pre_process.py

训练

$ python train.py

可视化训练过程,执行:

$ tensorboard --logdir path_to_current_dir/logs

演示

下载 预训练模型 放在 models 目录,然后执行:

$ python demo.py
原图 注意力
image
image
image
image
image
image
image
image
image
image

小小的赞助~

Sample

若对您有帮助可给予小小的赞助~