[NLP] 自己动手跑Google的Image Caption模型

两个月前Google公开了其之前在MSCOCO2015 Image Caption竞赛上夺得第一的Show&Tell模型（与微软MSR基于DSSM的模型并列）基于TensorFlow的实现，最近在做这方面的工作，就试着跑了一下。代码工程在gitub上。RNN和LSTM的一些基本情况介绍可以参看这里：[NL系列] RNN & LSTM 网络结构及应用。

Show&Tell/ im2txt

Google把公开之后的模型名称取为更像个工程名字的im2txt，其框架就像这张图：

图中 <code>S_{0}</code>到 <code>S_{N}</code>为生成的句子（包括开头和结尾各一个标识符），<code>W_{e}S_{i}</code> 为第 i 个词对应的词向量，LSTM 的输出 <code>p_{i}</code> 是模型生成的句子中下一个单词（第 i 个）的概率分布。<code>log p_{i}(S_{i})</code> 代表位置 i 生成的单词正确性的log-likelihoods，这些值的总和的负数就是模型的最小化目标。

Google的模型采用了End-to-end的思路，借用了机器翻译中的Encoder-Decoder框架（或者说是Google自己的Seq2Seq），通过一个模型直接将图像转换到句子。

机器翻译中Encoder-Decoder (Seq2Seq)模型的想法是，使用一个Encoder RNN读取源语言的句子，将其变换到一个固定长度的向量表示，然后使用 Decoder RNN将向量表示作为隐层初始值，产生目标语言的句子。

而im2txt的想法是，利用CNN在图片特征提取方面的强大能力，将Encoder RNN替换成CNN（im2txt中使用的是Google自己的Inception v3，模型在 ImageNet 分类任务上的准确率达到 93.9%，使得生成的图片描述的 BLEU-4 指标增加了 2 分），先利用CNN将图片转换到一个向量表示，再利用RNN将其转换到句子描述（采用beam search的方式，即迭代的在时刻t时保存k条最佳的句子片段用于生成t+1时刻的词，生成t+1时刻的词之后也只保存t+1时刻的k条最佳句子片段。代码中k选择的是3，论文中说的是20，应该是照顾了人民群众的基础设施肯定不如Google的关系）。

在实现中，im2txt基于在ILSVRC-2012-CLS 图片分类数据集上预训练好的CNN image recognition模型Inception v3，将其最后一个隐藏层作为Encoder RNN的输入，从而产生句子描述。

Before Preparation

虽然Google公开了其源码，但是想要自己训练一个im2txt模型并不是件容易的事，首先你得有一个能力足够的、可以运行CUDA的GPU。根据作者提供的信息，在一个NVIDIA Tesla K20m上进行初始训练大概需要1～2周的时间，如果为了达到更好的效果去进行fine tune的话，还需要再多几周才能达到peak performanc（应该就是论文中的数据）。虽然随时终止训练过程也可以得到效果不错的模型，但如果是要发论文刷分的话就得花不少时间。

Whilst it is possible to run this code on a CPU, beware that this may be approximately 10 times slower.

如果这些最基础的条件能够满足的话，就可以开始接下来的工作了。

Preparation

按照链接给出的教程，依次安装以下工具：

Bazel：Bazel是Google开源的自动化构建工具，类似于Make的功能，用来编译构建tensorflow。链接中给出的是Bazel官方在Ubuntu14.04或15.04下的安装教程，如果使用Java7的话可以按照这里的介绍稍作修改。
TensorFlow：注意安装的时候选择从源码编译的选项，按照支持GPU的步骤安装（所以首先要安装CUDA和CuDNN等）。中文版的教程可能在命令的版本上有所区别，最新版的建议看英文官网。
NumPy：基本上安装TensorFlow的时候都会装好。
Natural Language Toolkit (NLTK)：用于NLP的开源python函数库。首先安装NLTK，然后安装NLTK data.

安装完这些工具之后，就可以开始执行im2txt的自带脚本来下载需要的数据集，由于数据解压缩之后总共大概需要150GB的磁盘空间，因此建议先看看硬盘容量够不够，确定完之后执行以下命令：

# Location to save the MSCOCO data.
MSCOCO_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/mscoco"

# Build the preprocessing script.
bazel build im2txt/download_and_preprocess_mscoco

# Run the preprocessing script.
bazel-bin/im2txt/download_and_preprocess_mscoco "${MSCOCO_DIR}"

等到输出下面这句话，数据集的准备就算完成一半了。

2016-09-01 16:47:47.296630: Finished processing all 20267 image-caption pairs in data set 'test'.

剩下的一半是要把在ILSVRC-2012-CLS 图片分类数据集上预训练好的Inception v3模型下载下来。

This checkpoint file is provided by the TensorFlow-Slim image classification library which provides a suite of pre-trained image classification models.

执行以下命令（注意可以到TensorFlow-Slim image classification library看看最新的模型是什么，替换下面的 inception_v3_2016_08_28.tar.gz ：

# Location to save the Inception v3 checkpoint.
INCEPTION_DIR="${HOME}/im2txt/data"
mkdir -p ${INCEPTION_DIR}

wget "http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz"
tar -xvf "inception_v3_2016_08_28.tar.gz" -C ${INCEPTION_DIR}
rm "inception_v3_2016_08_28.tar.gz"

这个pre-trained模型只会在第一次执行训练时用到，im2txt每训练一段时间（默认的应该是迭代1024次）就会保存一次模型的checkpoint，之后的训练过程都会从checkpoint开始。

Start Training

im2txt的模型训练分为两步，第一步的initial training会固定CNN部分（Inception V3）的参数，把其当作一个图像编码网络生成image embedding，参与训练的只有在Inception V3上增加的一层网络（用于将image embedding映射到LSTM的word embedding vector space），而LSTM部分的所有待训练参数在此都会参与训练。

# Directory containing preprocessed MSCOCO data.
MSCOCO_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/mscoco"

# Inception v3 checkpoint file.
INCEPTION_CHECKPOINT="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/inception_v3.ckpt"

# Directory to save the model.
MODEL_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/model"

# Build the model.
bazel build -c opt im2txt/...

# Run the training script.
bazel-bin/im2txt/train \
  --input_file_pattern="${MSCOCO_DIR}/train-?????-of-00256" \
  --inception_checkpoint_file="${INCEPTION_CHECKPOINT}" \
  --train_dir="${MODEL_DIR}/train" \
  --train_inception=false \
  --number_of_steps=1000000

在训练的同时可以执行evaluation，以在TensorFlow自带的TensorBoard上方便的查看当前训练情况。如果只有一个GPU的话没有办法同时在GPU上跑evaluation（内存不够），因此一般是在CPU上执行，可以在命令行中执行export CUDA_VISIBLE_DEVICES=""命令限制当前程序看不到CUDA设备。默认的evaluation每600秒执行一次，在从最初的Inception V3模型迭代5000次之后才会开始，这些参数和设置都可以通过查看evaluate.py的代码了解。

MSCOCO_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/mscoco"
MODEL_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/model"

# Ignore GPU devices (only necessary if your GPU is currently memory
# constrained, for example, by running the training script).
export CUDA_VISIBLE_DEVICES=""

# Run the evaluation script. This will run in a loop, periodically loading the
# latest model checkpoint file and computing evaluation metrics.
bazel-bin/im2txt/evaluate \
  --input_file_pattern="${MSCOCO_DIR}/val-?????-of-00004" \
  --checkpoint_dir="${MODEL_DIR}/train" \
  --eval_dir="${MODEL_DIR}/eval"

然后就可以开启一个TensorBoard进程通过浏览器监控训练进度。

MODEL_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/model"

# Run a TensorBoard server.
tensorboard --logdir="${MODEL_DIR}"

Generating Captions

其实在训练的过程中随时可以生成图片描述，只是效果并不好说（其实也不一定比迭代很久之后差！）
执行以下命令：

# Directory containing model checkpoints.
CHECKPOINT_DIR="${YOUR_ADDR_TO_IM2TXT}/im2txt/model/train"

# Vocabulary file generated by the preprocessing script.
VOCAB_FILE="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/mscoco/word_counts.txt"

# JPEG image file to caption.
IMAGE_FILE="${YOUR_ADDR_TO_IM2TXT}/im2txt/data/mscoco/raw-data/val2014/${CHOICE_OF_IMAGE_}.jpg"

# Build the inference binary.
bazel build -c opt im2txt/run_inference

# Ignore GPU devices (only necessary if your GPU is currently memory
# constrained, for example, by running the training script).
export CUDA_VISIBLE_DEVICES=""

# Run inference to generate captions.
bazel-bin/im2txt/run_inference \
  --checkpoint_path=${CHECKPOINT_DIR} \
  --vocab_file=${VOCAB_FILE} \
  --input_files=${IMAGE_FILE}

官方给出的sample如下：

COCO_val2014_000000224477.jpg: a man riding a wave on top of a surfboard .

Captions for image COCO_val2014_000000224477.jpg:
  0) a man riding a wave on top of a surfboard . (p=0.040413)
  1) a person riding a surf board on a wave (p=0.017452)
  2) a man riding a wave on a surfboard in the ocean . (p=0.005743)

其实在我跑的时候大概迭代到200000次时同样是这张图生成的caption感觉比现在的第一条还要更合理一些，这就见仁见智了。

If you want more

如果之前的训练你觉得已经足够久，或者生成的caption你觉得还需要进一步优化，或者你正在苦于怎么超过state-of-art，那就可以把CNN的参数也一起放进来训练了，执行以下命令：

# Restart the training script with --train_inception=true.
bazel-bin/im2txt/train \
  --input_file_pattern="${MSCOCO_DIR}/train-?????-of-00256" \
  --train_dir="${MODEL_DIR}/train" \
  --train_inception=true \
  --number_of_steps=3000000  # Additional 2M steps (assuming 1M in initial training).

来自Google的温馨提醒：

Note that training will proceed much slower now, and the model will continue to improve by a small amount for a long time. We have found that it will improve slowly for an additional 2-2.5 million steps before it begins to overfit. This may take several weeks on a single GPU.

A Little Thoughts

还能有什么感想呢，现在initial training都没跑完。
这里有个日文的report和本文内容差不多，可以参考，里面有从TensorBoard中截取出来的图像。