ビジュアルアテンションを用いた画像キャプショニング

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

下記のような画像をもとに、"a surfer riding on a wave" のようなキャプションを生成することをめざします。

Man Surfing

画像ソース; ライセンス: パブリックドメイン

これを達成するため、アテンションベースのモデルを用います。これにより、キャプションを生成する際にモデルが画像のどの部分に焦点を当てているかを見ることができます。

Prediction

モデルの構造は、Show, Attend and Tell: Neural Image Caption Generation with Visual Attention と類似のものです。

このノートブックには例として、一連の処理の最初から最後までが含まれています。このノートブックを実行すると、MS-COCO データセットをダウンロードし、Inception V3 を使って画像のサブセットを前処理し、キャッシュします。その後、エンコーダー・デコーダーモデルを訓練し、訓練したモデルを使って新しい画像のキャプションを生成します。

この例では、比較的少量のデータ、およそ 20,000 枚の画像に対する最初の 30,000 のキャプションを使ってモデルを訓練します(データセットには 1 枚の画像あたり複数のキャプションがあるからです)。

import tensorflow as tf

# モデルがキャプション生成中に画像のどの部分に注目しているかを見るために
# アテンションのプロットを生成
import matplotlib.pyplot as plt

# Scikit-learn には役に立つさまざまなユーティリティが含まれる
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

import re
import numpy as np
import os
import time
import json
from glob import glob
from PIL import Image
import pickle

MS-COCO データセットのダウンロードと準備

モデルの訓練には MS-COCO データセット を使用します。このデータセットには、82,000 枚以上の画像が含まれ、それぞれの画像には少なくとも 5 つの異なったキャプションがつけられています。下記のコードは自動的にデータセットをダウンロードして解凍します。

annotation_zip = tf.keras.utils.get_file('captions.zip',
                                          cache_subdir=os.path.abspath('.'),
                                          origin = 'http://images.cocodataset.org/annotations/annotations_trainval2014.zip',
                                          extract = True)
annotation_file = os.path.dirname(annotation_zip)+'/annotations/captions_train2014.json'

name_of_zip = 'train2014.zip'
if not os.path.exists(os.path.abspath('.') + '/' + name_of_zip):
  image_zip = tf.keras.utils.get_file(name_of_zip,
                                      cache_subdir=os.path.abspath('.'),
                                      origin = 'http://images.cocodataset.org/zips/train2014.zip',
                                      extract = True)
  PATH = os.path.dirname(image_zip)+'/train2014/'
else:
  PATH = os.path.abspath('.')+'/train2014/'
Downloading data from http://images.cocodataset.org/annotations/annotations_trainval2014.zip
252878848/252872794 [==============================] - 20s 0us/step
Downloading data from http://images.cocodataset.org/zips/train2014.zip
13510574080/13510573713 [==============================] - 793s 0us/step

オプション: 訓練用データセットのサイズ制限

このチュートリアルの訓練をスピードアップするため、サブセットである 30,000 のキャプションと対応する画像を使ってモデルを訓練します。より多くのデータを使えばキャプションの品質が向上するでしょう。

# json ファイルの読み込み
with open(annotation_file, 'r') as f:
    annotations = json.load(f)

# ベクトルにキャプションと画像の名前を格納
all_captions = []
all_img_name_vector = []

for annot in annotations['annotations']:
    caption = '<start> ' + annot['caption'] + ' <end>'
    image_id = annot['image_id']
    full_coco_image_path = PATH + 'COCO_train2014_' + '%012d.jpg' % (image_id)

    all_img_name_vector.append(full_coco_image_path)
    all_captions.append(caption)

# captions と image_names を一緒にシャッフル
# random_state を設定
train_captions, img_name_vector = shuffle(all_captions,
                                          all_img_name_vector,
                                          random_state=1)

# シャッフルしたデータセットから最初の 30,000 のキャプションを選択
num_examples = 30000
train_captions = train_captions[:num_examples]
img_name_vector = img_name_vector[:num_examples]
len(train_captions), len(all_captions)
(30000, 414113)

InceptionV3 を使った画像の前処理

つぎに、(Imagenet を使って訓練済みの)InceptionV3をつかって、画像を分類します。最後の畳み込み層から特徴量を抽出します。

最初に、つぎのようにして画像を InceptionV3 が期待するフォーマットに変換します。

  • 画像を 299 ピクセル × 299 ピクセルにリサイズ
  • preprocess_input メソッドをつかって画像を InceptionV3 の訓練用画像のフォーマットに合致した −1 から 1 の範囲のピクセルを持つ形式に標準化して画像の前処理 を行う
def load_image(image_path):
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (299, 299))
    img = tf.keras.applications.inception_v3.preprocess_input(img)
    return img, image_path

InceptionV3 を初期化し Imagenet で学習済みの重みをロードする

ここで、出力層が InceptionV3 アーキテクチャの最後の畳み込み層である tf.keras モデルを作成しましょう。このレイヤーの出力の shape は8x8x2048 です。

  • 画像を 1 枚ずつネットワークに送り込み、処理結果のベクトルをディクショナリに保管します (image_name --> feature_vector)
  • すべての画像をネットワークで処理したあと、ディクショナリを pickle 形式でディスクに書き出します
image_model = tf.keras.applications.InceptionV3(include_top=False,
                                                weights='imagenet')
new_input = image_model.input
hidden_layer = image_model.layers[-1].output

image_features_extract_model = tf.keras.Model(new_input, hidden_layer)
Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/inception_v3/inception_v3_weights_tf_dim_ordering_tf_kernels_notop.h5
87916544/87910968 [==============================] - 4s 0us/step

InceptionV3 から抽出した特徴量のキャッシング

それぞれの画像を InceptionV3 で前処理して出力をディスクにキャッシュします。出力を RAM にキャッシュすれば速くなりますが、画像 1 枚あたり 8 * 8 * 2048 個の浮動小数点数が必要で、メモリを大量に必要とします。これを書いている時点では、これは Colab のメモリ上限(現在は 12GB)を超えています。

より高度なキャッシング戦略(たとえば画像を分散保存しディスクのランダムアクセス入出力を低減するなど)を使えば性能は向上しますが、より多くのコーディングが必要となります。

このキャッシングは Colab で GPU を使った場合で約 10 分ほどかかります。プログレスバーを表示したければ次のようにします。

  1. tqdm をインストールします:

    !pip install -q tqdm

  2. tqdm をインポートします:

    from tqdm import tqdm

  3. 下記の行を:

    for img, path in image_dataset:

    次のように変更します:

    for img, path in tqdm(image_dataset):

# 重複のない画像を取得
encode_train = sorted(set(img_name_vector))

# batch_size はシステム構成に合わせて自由に変更可能
image_dataset = tf.data.Dataset.from_tensor_slices(encode_train)
image_dataset = image_dataset.map(
  load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(16)

for img, path in image_dataset:
  batch_features = image_features_extract_model(img)
  batch_features = tf.reshape(batch_features,
                              (batch_features.shape[0], -1, batch_features.shape[3]))

  for bf, p in zip(batch_features, path):
    path_of_feature = p.numpy().decode("utf-8")
    np.save(path_of_feature, bf.numpy())

キャプションの前処理とトークン化

  • まず最初に、キャプションをトークン化(たとえばスペースで区切るなど)します。これにより、データ中の重複しない単語のボキャブラリ(たとえば、"surfing"、"football" など)が得られます。
  • つぎに、(メモリ節約のため)ボキャブラリのサイズを上位 5,000 語に制限します。それ以外の単語は "UNK" (不明)というトークンに置き換えます。
  • 続いて、単語からインデックス、インデックスから単語への対応表を作成します。
  • 最後に、すべてのシーケンスの長さを最長のものに合わせてパディングします。
# データセット中の一番長いキャプションの長さを検出
def calc_max_length(tensor):
    return max(len(t) for t in tensor)
# ボキャブラリ中のトップ 5000 語を選択
top_k = 5000
tokenizer = tf.keras.preprocessing.text.Tokenizer(num_words=top_k,
                                                  oov_token="<unk>",
                                                  filters='!"#$%&()*+.,-/:;=?@[\]^_`{|}~ ')
tokenizer.fit_on_texts(train_captions)
train_seqs = tokenizer.texts_to_sequences(train_captions)
tokenizer.word_index['<pad>'] = 0
tokenizer.index_word[0] = '<pad>'
# トークン化したベクトルを生成
train_seqs = tokenizer.texts_to_sequences(train_captions)
# キャプションの最大長に各ベクトルをパディング
# max_length を指定しない場合、pad_sequences は自動的に計算
cap_vector = tf.keras.preprocessing.sequence.pad_sequences(train_seqs, padding='post')
# アテンションの重みを格納するために使われる max_length を計算
max_length = calc_max_length(train_seqs)

データを訓練用とテスト用に分割

# 訓練用セットと検証用セットを 80-20 に分割して生成
img_name_train, img_name_val, cap_train, cap_val = train_test_split(img_name_vector,
                                                                    cap_vector,
                                                                    test_size=0.2,
                                                                    random_state=0)
len(img_name_train), len(cap_train), len(img_name_val), len(cap_val)
(24000, 24000, 6000, 6000)

訓練用の tf.data データセットの作成

画像とキャプションが用意できました。つぎは、モデルの訓練に使用する tf.data データセットを作成しましょう。

# これらのパラメータはシステム構成に合わせて自由に変更してください

BATCH_SIZE = 64
BUFFER_SIZE = 1000
embedding_dim = 256
units = 512
vocab_size = len(tokenizer.word_index) + 1
num_steps = len(img_name_train) // BATCH_SIZE
# InceptionV3 から抽出したベクトルの shape は (64, 2048)
# つぎの 2 つのパラメータはこのベクトルの shape を表す
features_shape = 2048
attention_features_shape = 64
# numpy ファイルをロード
def map_func(img_name, cap):
  img_tensor = np.load(img_name.decode('utf-8')+'.npy')
  return img_tensor, cap
dataset = tf.data.Dataset.from_tensor_slices((img_name_train, cap_train))

# numpy ファイルを並列に読み込むために map を使用
dataset = dataset.map(lambda item1, item2: tf.numpy_function(
          map_func, [item1, item2], [tf.float32, tf.int32]),
          num_parallel_calls=tf.data.experimental.AUTOTUNE)

# シャッフルとバッチ化
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

モデル

興味深い事実:下記のデコーダはNeural Machine Translation with Attention の例のデコーダとまったく同一です。

このモデルのアーキテクチャは、Show, Attend and Tell の論文にインスパイアされたものです。

  • この例では、InceptionV3 の下層の畳込みレイヤーから特徴量を取り出します。得られるベクトルの shape は (8, 8, 2048) です。
  • このベクトルを (64, 2048) に変形します。
  • このベクトルは(1層の全結合層からなる)CNN エンコーダに渡されます。
  • RNN(ここではGRU)が画像を介して次の単語を予測します。
class BahdanauAttention(tf.keras.Model):
  def __init__(self, units):
    super(BahdanauAttention, self).__init__()
    self.W1 = tf.keras.layers.Dense(units)
    self.W2 = tf.keras.layers.Dense(units)
    self.V = tf.keras.layers.Dense(1)

  def call(self, features, hidden):
    # features(CNN_encoder output) shape == (batch_size, 64, embedding_dim)

    # hidden shape == (batch_size, hidden_size)
    # hidden_with_time_axis shape == (batch_size, 1, hidden_size)
    hidden_with_time_axis = tf.expand_dims(hidden, 1)

    # score shape == (batch_size, 64, hidden_size)
    score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))

    # attention_weights shape == (batch_size, 64, 1)
    # score を self.V に適用するので、最後の軸は 1 となる
    attention_weights = tf.nn.softmax(self.V(score), axis=1)

    # 合計をとったあとの context_vector の shpae == (batch_size, hidden_size)
    context_vector = attention_weights * features
    context_vector = tf.reduce_sum(context_vector, axis=1)

    return context_vector, attention_weights
class CNN_Encoder(tf.keras.Model):
    # すでに特徴量を抽出して pickle 形式でダンプしてあるので
    # このエンコーダはそれらの特徴量を全結合層に渡して処理する
    def __init__(self, embedding_dim):
        super(CNN_Encoder, self).__init__()
        # shape after fc == (batch_size, 64, embedding_dim)
        self.fc = tf.keras.layers.Dense(embedding_dim)

    def call(self, x):
        x = self.fc(x)
        x = tf.nn.relu(x)
        return x
class RNN_Decoder(tf.keras.Model):
  def __init__(self, embedding_dim, units, vocab_size):
    super(RNN_Decoder, self).__init__()
    self.units = units

    self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
    self.gru = tf.keras.layers.GRU(self.units,
                                   return_sequences=True,
                                   return_state=True,
                                   recurrent_initializer='glorot_uniform')
    self.fc1 = tf.keras.layers.Dense(self.units)
    self.fc2 = tf.keras.layers.Dense(vocab_size)

    self.attention = BahdanauAttention(self.units)

  def call(self, x, features, hidden):
    # アテンションを別のモデルとして定義
    context_vector, attention_weights = self.attention(features, hidden)

    # embedding 層を通過したあとの x の shape == (batch_size, 1, embedding_dim)
    x = self.embedding(x)

    # 結合後の x の shape == (batch_size, 1, embedding_dim + hidden_size)
    x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)

    # 結合したベクトルを GRU に渡す
    output, state = self.gru(x)

    # shape == (batch_size, max_length, hidden_size)
    x = self.fc1(output)

    # x shape == (batch_size * max_length, hidden_size)
    x = tf.reshape(x, (-1, x.shape[2]))

    # output shape == (batch_size * max_length, vocab)
    x = self.fc2(x)

    return x, state, attention_weights

  def reset_state(self, batch_size):
    return tf.zeros((batch_size, self.units))
encoder = CNN_Encoder(embedding_dim)
decoder = RNN_Decoder(embedding_dim, units, vocab_size)
optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')

def loss_function(real, pred):
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  loss_ = loss_object(real, pred)

  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask

  return tf.reduce_mean(loss_)

チェックポイント

checkpoint_path = "./checkpoints/train"
ckpt = tf.train.Checkpoint(encoder=encoder,
                           decoder=decoder,
                           optimizer = optimizer)
ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)
start_epoch = 0
if ckpt_manager.latest_checkpoint:
  start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])

訓練

  • それぞれの .npy ファイルに保存されている特徴量を取り出し、エンコーダに渡す
  • エンコーダの出力と(0 で初期化された)隠れ状態と、デコーダの入力(開始トークン)をデコーダに渡す
  • デコーダは予測値とデコーダの隠れ状態を返す
  • デコーダの隠れ状態はモデルに戻され、予測値をつかって損失を計算する
  • デコーダの次の入力を決めるために teacher forcing を用いる
  • Teacher forcing は、デコーダの次の入力として正解の単語を渡す手法である
  • 最後のステップは、勾配を計算してそれをオプティマイザに適用し誤差逆伝播を行うことである
# 訓練用のセルを複数回実行すると loss_plot 配列がリセットされてしまうので
# 独立したセルとして追加
loss_plot = []
@tf.function
def train_step(img_tensor, target):
  loss = 0

  # バッチごとに隠れ状態を初期化
  # 画像のキャプションはその前後の画像と無関係なため
  hidden = decoder.reset_state(batch_size=target.shape[0])

  dec_input = tf.expand_dims([tokenizer.word_index['<start>']] * target.shape[0], 1)

  with tf.GradientTape() as tape:
      features = encoder(img_tensor)

      for i in range(1, target.shape[1]):
          # 特徴量をデコーダに渡す
          predictions, hidden, _ = decoder(dec_input, features, hidden)

          loss += loss_function(target[:, i], predictions)

          # teacher forcing を使用
          dec_input = tf.expand_dims(target[:, i], 1)

  total_loss = (loss / int(target.shape[1]))

  trainable_variables = encoder.trainable_variables + decoder.trainable_variables

  gradients = tape.gradient(loss, trainable_variables)

  optimizer.apply_gradients(zip(gradients, trainable_variables))

  return loss, total_loss
EPOCHS = 20

for epoch in range(start_epoch, EPOCHS):
    start = time.time()
    total_loss = 0

    for (batch, (img_tensor, target)) in enumerate(dataset):
        batch_loss, t_loss = train_step(img_tensor, target)
        total_loss += t_loss

        if batch % 100 == 0:
            print ('Epoch {} Batch {} Loss {:.4f}'.format(
              epoch + 1, batch, batch_loss.numpy() / int(target.shape[1])))
    # 後ほどグラフ化するためにエポックごとに損失を保存
    loss_plot.append(total_loss / num_steps)

    if epoch % 5 == 0:
      ckpt_manager.save()

    print ('Epoch {} Loss {:.6f}'.format(epoch + 1,
                                         total_loss/num_steps))
    print ('Time taken for 1 epoch {} sec\n'.format(time.time() - start))
Epoch 1 Batch 0 Loss 2.1039
Epoch 1 Batch 100 Loss 1.1519
Epoch 1 Batch 200 Loss 1.0072
Epoch 1 Batch 300 Loss 0.9093
Epoch 1 Loss 1.070060
Time taken for 1 epoch 80.72183585166931 sec

Epoch 2 Batch 0 Loss 0.7917
Epoch 2 Batch 100 Loss 0.9244
Epoch 2 Batch 200 Loss 0.7865
Epoch 2 Batch 300 Loss 0.8113
Epoch 2 Loss 0.804910
Time taken for 1 epoch 38.883573055267334 sec

Epoch 3 Batch 0 Loss 0.7552
Epoch 3 Batch 100 Loss 0.7392
Epoch 3 Batch 200 Loss 0.7135
Epoch 3 Batch 300 Loss 0.7433
Epoch 3 Loss 0.732453
Time taken for 1 epoch 38.73031163215637 sec

Epoch 4 Batch 0 Loss 0.6670
Epoch 4 Batch 100 Loss 0.7263
Epoch 4 Batch 200 Loss 0.6449
Epoch 4 Batch 300 Loss 0.6899
Epoch 4 Loss 0.685711
Time taken for 1 epoch 38.65649175643921 sec

Epoch 5 Batch 0 Loss 0.6861
Epoch 5 Batch 100 Loss 0.7060
Epoch 5 Batch 200 Loss 0.6696
Epoch 5 Batch 300 Loss 0.6174
Epoch 5 Loss 0.648918
Time taken for 1 epoch 38.62453603744507 sec

Epoch 6 Batch 0 Loss 0.6339
Epoch 6 Batch 100 Loss 0.6359
Epoch 6 Batch 200 Loss 0.6329
Epoch 6 Batch 300 Loss 0.6083
Epoch 6 Loss 0.616448
Time taken for 1 epoch 39.09868240356445 sec

Epoch 7 Batch 0 Loss 0.6181
Epoch 7 Batch 100 Loss 0.6222
Epoch 7 Batch 200 Loss 0.6163
Epoch 7 Batch 300 Loss 0.5793
Epoch 7 Loss 0.586188
Time taken for 1 epoch 38.76476716995239 sec

Epoch 8 Batch 0 Loss 0.5797
Epoch 8 Batch 100 Loss 0.5801
Epoch 8 Batch 200 Loss 0.5079
Epoch 8 Batch 300 Loss 0.5723
Epoch 8 Loss 0.557516
Time taken for 1 epoch 38.59176993370056 sec

Epoch 9 Batch 0 Loss 0.5109
Epoch 9 Batch 100 Loss 0.5736
Epoch 9 Batch 200 Loss 0.5730
Epoch 9 Batch 300 Loss 0.4985
Epoch 9 Loss 0.526951
Time taken for 1 epoch 38.505391120910645 sec

Epoch 10 Batch 0 Loss 0.4796
Epoch 10 Batch 100 Loss 0.5103
Epoch 10 Batch 200 Loss 0.4917
Epoch 10 Batch 300 Loss 0.4898
Epoch 10 Loss 0.497396
Time taken for 1 epoch 38.57945775985718 sec

Epoch 11 Batch 0 Loss 0.5007
Epoch 11 Batch 100 Loss 0.4797
Epoch 11 Batch 200 Loss 0.4580
Epoch 11 Batch 300 Loss 0.4676
Epoch 11 Loss 0.470451
Time taken for 1 epoch 38.64489150047302 sec

Epoch 12 Batch 0 Loss 0.4361
Epoch 12 Batch 100 Loss 0.4548
Epoch 12 Batch 200 Loss 0.4174
Epoch 12 Batch 300 Loss 0.4617
Epoch 12 Loss 0.441208
Time taken for 1 epoch 38.6077823638916 sec

Epoch 13 Batch 0 Loss 0.4561
Epoch 13 Batch 100 Loss 0.4436
Epoch 13 Batch 200 Loss 0.4144
Epoch 13 Batch 300 Loss 0.4070
Epoch 13 Loss 0.414397
Time taken for 1 epoch 38.550392389297485 sec

Epoch 14 Batch 0 Loss 0.4009
Epoch 14 Batch 100 Loss 0.3738
Epoch 14 Batch 200 Loss 0.3684
Epoch 14 Batch 300 Loss 0.4379
Epoch 14 Loss 0.388811
Time taken for 1 epoch 38.69609332084656 sec

Epoch 15 Batch 0 Loss 0.4075
Epoch 15 Batch 100 Loss 0.3639
Epoch 15 Batch 200 Loss 0.3469
Epoch 15 Batch 300 Loss 0.3569
Epoch 15 Loss 0.364227
Time taken for 1 epoch 38.454240560531616 sec

Epoch 16 Batch 0 Loss 0.3631
Epoch 16 Batch 100 Loss 0.3610
Epoch 16 Batch 200 Loss 0.3278
Epoch 16 Batch 300 Loss 0.3265
Epoch 16 Loss 0.339858
Time taken for 1 epoch 38.848779916763306 sec

Epoch 17 Batch 0 Loss 0.3299
Epoch 17 Batch 100 Loss 0.3294
Epoch 17 Batch 200 Loss 0.3136
Epoch 17 Batch 300 Loss 0.3095
Epoch 17 Loss 0.318829
Time taken for 1 epoch 38.28813695907593 sec

Epoch 18 Batch 0 Loss 0.3175
Epoch 18 Batch 100 Loss 0.3092
Epoch 18 Batch 200 Loss 0.2808
Epoch 18 Batch 300 Loss 0.3057
Epoch 18 Loss 0.300110
Time taken for 1 epoch 38.392847299575806 sec

Epoch 19 Batch 0 Loss 0.3309
Epoch 19 Batch 100 Loss 0.3000
Epoch 19 Batch 200 Loss 0.2641
Epoch 19 Batch 300 Loss 0.2619
Epoch 19 Loss 0.280883
Time taken for 1 epoch 38.31717014312744 sec

Epoch 20 Batch 0 Loss 0.3077
Epoch 20 Batch 100 Loss 0.2751
Epoch 20 Batch 200 Loss 0.2310
Epoch 20 Batch 300 Loss 0.2734
Epoch 20 Loss 0.261823
Time taken for 1 epoch 38.140281677246094 sec


plt.plot(loss_plot)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Loss Plot')
plt.show()

png

キャプション!

  • 評価関数は訓練ループとおなじだが、teacher forcing は使わない。タイムステップごとのデコーダへの入力は、隠れ状態とエンコーダの入力に加えて、一つ前の予測値である。
  • モデルが終了トークンを予測したら、予測を終了する。
  • それぞれのタイムステップごとに、アテンションの重みを保存する。
def evaluate(image):
    attention_plot = np.zeros((max_length, attention_features_shape))

    hidden = decoder.reset_state(batch_size=1)

    temp_input = tf.expand_dims(load_image(image)[0], 0)
    img_tensor_val = image_features_extract_model(temp_input)
    img_tensor_val = tf.reshape(img_tensor_val, (img_tensor_val.shape[0], -1, img_tensor_val.shape[3]))

    features = encoder(img_tensor_val)

    dec_input = tf.expand_dims([tokenizer.word_index['<start>']], 0)
    result = []

    for i in range(max_length):
        predictions, hidden, attention_weights = decoder(dec_input, features, hidden)

        attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()

        predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()
        result.append(tokenizer.index_word[predicted_id])

        if tokenizer.index_word[predicted_id] == '<end>':
            return result, attention_plot

        dec_input = tf.expand_dims([predicted_id], 0)

    attention_plot = attention_plot[:len(result), :]
    return result, attention_plot
def plot_attention(image, result, attention_plot):
    temp_image = np.array(Image.open(image))

    fig = plt.figure(figsize=(10, 10))

    len_result = len(result)
    for l in range(len_result):
        temp_att = np.resize(attention_plot[l], (8, 8))
        ax = fig.add_subplot(len_result//2, len_result//2, l+1)
        ax.set_title(result[l])
        img = ax.imshow(temp_image)
        ax.imshow(temp_att, cmap='gray', alpha=0.6, extent=img.get_extent())

    plt.tight_layout()
    plt.show()
# 検証用セットのキャプション
rid = np.random.randint(0, len(img_name_val))
image = img_name_val[rid]
real_caption = ' '.join([tokenizer.index_word[i] for i in cap_val[rid] if i not in [0]])
result, attention_plot = evaluate(image)

print ('Real Caption:', real_caption)
print ('Prediction Caption:', ' '.join(result))
plot_attention(image, result, attention_plot)

Real Caption: <start> there is a train on the tracks below all the wires <end>
Prediction Caption: a train are getting ready to board pulling a platform while tracks is parked while red train is traveling near tracks <end>

png

あなた独自の画像でためそう

お楽しみのために、訓練したばかりのモデルであなたの独自の画像を使うためのメソッドを下記に示します。比較的少量のデータで訓練していること、そして、あなたの画像は訓練データとは異なるであろうことを、心に留めておいてください(変な結果が出てくることを覚悟しておいてください)。

image_url = 'https://tensorflow.org/images/surf.jpg'
image_extension = image_url[-4:]
image_path = tf.keras.utils.get_file('image'+image_extension,
                                     origin=image_url)

result, attention_plot = evaluate(image_path)
print ('Prediction Caption:', ' '.join(result))
plot_attention(image_path, result, attention_plot)
# 画像を開く
Image.open(image_path)
Prediction Caption: two men look in a hat <end>

png

png

次のステップ

おめでとうございます!アテンション付きの画像キャプショニングモデルの訓練が終わりました。つぎは、この例 Neural Machine Translation with Attention を御覧ください。これはおなじアーキテクチャをつかってスペイン語と英語の間の翻訳を行います。このノートブックのコードをつかって、別のデータセットで訓練を行うこともできます。