Generate Image Caption

  依旧采用十分熟悉的NMT架构，把生成图像描述看作是图像到文本的翻译过程。

架构：

  模型的解码器部分主要由注意力层组成：因果自注意力，用于处理生成的文本序列。交叉注意力，用于处理生成文本和图像之间的注意力。

数据简介

数据来自一个非公开的中文数据集，每个图像有5个不同的中文描述 。

import os
import json
import jieba
import einops
import collections
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as pltfrom PIL import Image
from tqdm import tqdmos.environ["CUDA_VISIBLE_DEVICES"] = '1'os.environ['TF_CPP_MIN_LOG_LEVEL']='2' 

加载数据

with open('./caption_train/caption_train_annotations.json','r') as f:captions = json.loads(f.read())caption_df = pd.DataFrame(captions)del caption_df['url']

caption_df.head(3)

分词

caption_df['caption'] = caption_df.caption.apply(lambda x: [' '.join(jieba.cut(_)) for _ in x])

Building prefix dict from the default dictionary ...
Loading model from cache /tmp/jieba.cache
Loading model cost 0.771 seconds.
Prefix dict has been built successfully.

caption_df.head()

划分训练测试集

test_df = caption_df.sample(10000)train_image_ids = set(caption_df.image_id) - set(test_df.image_id)train_df = caption_df[caption_df.image_id.isin(train_image_ids)]

train_raw = tf.data.Dataset.from_tensor_slices((train_df.image_id, train_df.caption.tolist()))test_raw = tf.data.Dataset.from_tensor_slices((test_df.image_id, test_df.caption.tolist()))

for ex_path, ex_captions in train_raw.take(1):print(ex_path)print(ex_captions.numpy()[0].decode())print(ex_captions.numpy()[1].decode())print(ex_captions.numpy()[2].decode())print(ex_captions.numpy()[3].decode())print(ex_captions.numpy()[4].decode())

tf.Tensor(b'8f00f3d0f1008e085ab660e70dffced16a8259f6.jpg', shape=(), dtype=string)
两个 衣着 休闲 的 人 在 平整 的 道路 上 交谈
一个 穿着 红色 上衣 的 男人 和 一个 穿着 灰色 裤子 的 男人 站 在 室外 的 道路 上 交谈
室外 的 公园 里 有 两个 穿着 长裤 的 男人 在 交流
街道 上 有 一个 穿着 深色 外套 的 男人 和 一个 穿着 红色 外套 的 男人 在 交谈
道路 上 有 一个 身穿 红色 上衣 的 男人 在 和 一个 抬着 左手 的 人 讲话

数据预处理

• 加载和缩放图像：缩放图像为MobileNet输入的大小
• 文本向量化：word→\to→token，token→\to→word
• 数据对齐：一个图对应5个中文描述，转换为1对1

# 加载和缩放图像
IMAGE_SHAPE = (224, 224, 3)
def load_image(image_name):img = tf.io.read_file('./caption_train/caption_train_images/'+image_name)img = tf.io.decode_jpeg(img, channels=3)img = tf.image.resize(img, IMAGE_SHAPE[:-1])/255.0return img

def load_test_image(image_path):img = tf.io.read_file(image_path)img = tf.io.decode_jpeg(img, channels=3)img = tf.image.resize(img, IMAGE_SHAPE[:-1])return img

test_img_batch = load_image(ex_path)[tf.newaxis, :]print(test_img_batch.shape)

(1, 224, 224, 3)

文本描述转换为token序列：

# 添加开始，和结束字符
def standardize(s):s = tf.strings.join(['[START]', s, '[END]'], separator=' ')return s

# Use the top 20000 words for a vocabulary.
vocabulary_size = 20000
tokenizer = tf.keras.layers.TextVectorization(max_tokens=vocabulary_size,standardize=standardize,ragged=True)# Learn the vocabulary from the caption data.
tokenizer.adapt(train_raw.map(lambda fname,caption: caption).unbatch().batch(1024))

tokenizer.get_vocabulary()[:10]

['', '[UNK]', '的', '[START]', '[END]', '一个', '在', '上', '男人', '着']

for i in range(5):print(ex_captions.numpy()[i].decode())
tokenizer(ex_captions)

两个 衣着 休闲 的 人 在 平整 的 道路 上 交谈
一个 穿着 红色 上衣 的 男人 和 一个 穿着 灰色 裤子 的 男人 站 在 室外 的 道路 上 交谈
室外 的 公园 里 有 两个 穿着 长裤 的 男人 在 交流
街道 上 有 一个 穿着 深色 外套 的 男人 和 一个 穿着 红色 外套 的 男人 在 交谈
道路 上 有 一个 身穿 红色 上衣 的 男人 在 和 一个 抬着 左手 的 人 讲话

vocab_size = len(tokenizer.get_vocabulary())
vocab_size

17303

数据对齐：
一个图片对应5条中文描述，对齐为1：1。

def match_shapes(images, captions):caption_shape = einops.parse_shape(captions, 'b c')captions = einops.rearrange(captions, 'b c -> (b c)')images = einops.repeat(images, 'b ... -> (b c) ...',c = caption_shape['c'])return images, captions

for ex_paths, ex_captions in train_raw.batch(32).take(1):breakprint("对齐前：")
print('image paths:', ex_paths.shape)
print('captions:', ex_captions.shape)
print()ex_paths, ex_captions = match_shapes(images=ex_paths, captions=ex_captions)
print("对齐后：")
print('image_paths:', ex_paths.shape)
print('captions:', ex_captions.shape)

对齐前：
image paths: (32,)
captions: (32, 5)对齐后：
image_paths: (160,)
captions: (160,)

训练数据格式（inputs, labels)：

• inputs : (images, input_tokens)
• labels : label_tokens

input_tokens和output_tokens都是通过图像的描述文本(token)序列得到，例如：

• tokens : [[1, 2, 3, 4, 5, 6, 7, 8]]
• input_tokens : [[1, 2, 3, 4, 5, 6, 7]]
• label_tokens : [[2, 3, 4, 5, 6, 7, 8]]

label_tokens的每一个位置的token正是input_tokes每个位置的下一个token。

def prepare_txt(imgs, txts):tokens = tokenizer(txts)input_tokens = tokens[..., :-1]label_tokens = tokens[..., 1:]return (imgs, input_tokens), label_tokens

def prepare_dataset(ds, tokenizer, batch_size=32, shuffle_buffer=1000):# Load the images and make batches.ds = (ds.shuffle(26000).map(lambda path, caption: (load_image(path), caption)).apply(tf.data.Dataset.ignore_errors).batch(batch_size))def to_tensor(inputs, labels):(images, in_tok), out_tok = inputs, labelsreturn (images, in_tok.to_tensor()), out_tok.to_tensor()return (ds.map(match_shapes, tf.data.AUTOTUNE).unbatch().shuffle(shuffle_buffer).batch(batch_size).map(prepare_txt, tf.data.AUTOTUNE).map(to_tensor, tf.data.AUTOTUNE))

train_ds = prepare_dataset(train_raw, tokenizer)test_ds = prepare_dataset(test_raw, tokenizer)
test_ds.element_spec

((TensorSpec(shape=(None, 224, 224, 3), dtype=tf.float32, name=None),TensorSpec(shape=(None, None), dtype=tf.int64, name=None)),TensorSpec(shape=(None, None), dtype=tf.int64, name=None))

模型组件

模型由四部分组成：

1. MobileNet 图像特征提取
2. Embedding 词和位置的嵌入
3. Transformer 解码层，每一层包含三个子层：
   • 因果自注意力层
   • 交叉注意力层
   • 前馈神经网络层
4. 输出层，预测下一个token的概率分布

特征提取

预训练模型：MobileNet

# 加载MobileNet 
IMAGE_SHAPE=(224, 224, 3)
mobilenet = tf.keras.applications.MobileNet(weights='/home/wjh/keras_weights/mobilenet_1_0_224_tf_no_top.h5',input_shape=IMAGE_SHAPE,include_top=False)
mobilenet.trainable=False

# 剪裁mobileNet
#pretrain_model = tf.keras.models.Model(inputs=mobilenet.input, outputs=mobilenet.get_layer('conv_pw_11_relu').output)print(mobilenet(test_img_batch).shape)

(1, 7, 7, 1024)

嵌入层

与CNN或RNN不同，注意力层对序列的顺序是不敏感的。如果没有一些位置输入，它只看到一个无序集合而不是一个序列。

class SeqEmbedding(tf.keras.layers.Layer):def __init__(self, vocab_size, max_length, embed_dim):super().__init__()self.pos_embedding = tf.keras.layers.Embedding(input_dim=max_length, output_dim=embed_dim)self.token_embedding = tf.keras.layers.Embedding(input_dim=vocab_size,output_dim=embed_dim,mask_zero=True)self.add = tf.keras.layers.Add()def call(self, seq):seq = self.token_embedding(seq) # (batch, seq, embed_dim)x = tf.range(tf.shape(seq)[1])  # (seq)x = x[tf.newaxis, :]            # (1, seq)x = self.pos_embedding(x)       # (1, seq, embed_dim)return self.add([seq,x])

因果自注意力层

因果自注意力跟自注意力的区别就是增加掩码，屏蔽当前token所在位置以后的tokens。

class CausalSelfAttention(tf.keras.layers.Layer):def __init__(self, **kwargs):super().__init__()self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)# Use Add instead of + so the keras mask propagates through.self.add = tf.keras.layers.Add() self.layernorm = tf.keras.layers.LayerNormalization()def call(self, x):attn = self.mha(query=x, value=x, use_causal_mask=True)x = self.add([x, attn])return self.layernorm(x)

交叉注意力层

交叉注意力用于处理图像和文本之间的交叉注意力，通过观察交叉注意力层的注意力权重，可以看到模型在生成单词时的注意力在图像的哪些部分。

class CrossAttention(tf.keras.layers.Layer):def __init__(self,**kwargs):super().__init__()self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)self.add = tf.keras.layers.Add() self.layernorm = tf.keras.layers.LayerNormalization()def call(self, x, y, **kwargs):attn, attention_scores = self.mha(query=x, value=y, return_attention_scores=True)self.last_attention_scores = attention_scoresx = self.add([x, attn])return self.layernorm(x)

前馈全连接层

全连接层，可能就是为了纯粹增加模型的参数，该层的作用是把张量的最后一个维度投射到高维空间然后在压缩到输入的维度。

class FeedForward(tf.keras.layers.Layer):def __init__(self, units, dropout_rate=0.1):super().__init__()self.seq = tf.keras.Sequential([tf.keras.layers.Dense(units=2*units, activation='relu'),tf.keras.layers.Dense(units=units),tf.keras.layers.Dropout(rate=dropout_rate),])self.layernorm = tf.keras.layers.LayerNormalization()def call(self, x):x = x + self.seq(x)return self.layernorm(x)

解码层

把因果自注意力层+交叉注意力层+前馈全连接层组装在一起：

class DecoderLayer(tf.keras.layers.Layer):def __init__(self, units, num_heads=1, dropout_rate=0.1):super().__init__()self.self_attention = CausalSelfAttention(num_heads=num_heads,key_dim=units,dropout=dropout_rate)self.cross_attention = CrossAttention(num_heads=num_heads,key_dim=units,dropout=dropout_rate)self.ff = FeedForward(units=units, dropout_rate=dropout_rate)def call(self, inputs, training=False):in_seq, out_seq = inputs# Text inputout_seq = self.self_attention(out_seq)out_seq = self.cross_attention(out_seq, in_seq)self.last_attention_scores = self.cross_attention.last_attention_scoresout_seq = self.ff(out_seq)return out_seq

输出层

预测下个token的分布，为获得更好的输出，可以通过一下两点来实现：

1. 处理无意义的字符：例如：'', '[UNK]', '[START]'在输出层中把偏置设置较大的负值，避免这些字符生成。
2. 最优初始偏差：统计token的数量，设置初始偏差为-p*log(p)。

class TokenOutput(tf.keras.layers.Layer):def __init__(self, tokenizer, banned_tokens=('', '[UNK]', '[START]'), **kwargs):super().__init__()self.dense = tf.keras.layers.Dense(units=tokenizer.vocabulary_size(), **kwargs)self.tokenizer = tokenizerself.banned_tokens = banned_tokensself.bias = Nonedef adapt(self, ds):counts = collections.Counter()vocab_dict = {name: id for id, name in enumerate(self.tokenizer.get_vocabulary())}for tokens in tqdm(ds):counts.update(tokens.numpy().flatten())counts_arr = np.zeros(shape=(self.tokenizer.vocabulary_size(),))counts_arr[np.array(list(counts.keys()), dtype=np.int32)] = list(counts.values())counts_arr = counts_arr[:]for token in self.banned_tokens:counts_arr[vocab_dict[token]] = 0total = counts_arr.sum()p = counts_arr/totalp[counts_arr==0] = 1.0log_p = np.log(p)  # log(1) == 0entropy = -(log_p*p).sum()print()print(f"Uniform entropy: {np.log(self.tokenizer.vocabulary_size()):0.2f}")print(f"Marginal entropy: {entropy:0.2f}")self.bias = log_pself.bias[counts_arr==0] = -1e9def call(self, x):x = self.dense(x)return x + self.bias

output_layer = TokenOutput(tokenizer, banned_tokens=('', '[UNK]', '[START]'))
# This might run a little faster if the dataset didn't also have to load the image data.
output_layer.adapt(train_ds.map(lambda inputs, labels: labels))

31250it [05:47, 89.89it/s] 

​
Uniform entropy: 9.76
Marginal entropy: 4.63

​

构建模型

class Captioner(tf.keras.Model):@classmethoddef add_method(cls, fun):setattr(cls, fun.__name__, fun)return fundef __init__(self, tokenizer, feature_extractor, output_layer, num_layers=1,units=256, max_length=50, num_heads=1, dropout_rate=0.1):super().__init__()self.feature_extractor = feature_extractorself.tokenizer = tokenizerself.word_to_index = tf.keras.layers.StringLookup(mask_token="",vocabulary=tokenizer.get_vocabulary())self.index_to_word = tf.keras.layers.StringLookup(mask_token="",vocabulary=tokenizer.get_vocabulary(),invert=True) self.seq_embedding = SeqEmbedding(vocab_size=tokenizer.vocabulary_size(),embed_dim=units,max_length=max_length)self.decoder_layers = [DecoderLayer(units, num_heads=num_heads, dropout_rate=dropout_rate)for n in range(num_layers)]self.output_layer = output_layer

@Captioner.add_method
def call(self, inputs):image, txt = inputsif image.shape[-1] == 3:# Apply the feature-extractor, if you get an RGB image.image = self.feature_extractor(image)# Flatten the feature mapimage = einops.rearrange(image, 'b h w c -> b (h w) c')if txt.dtype == tf.string:# Apply the tokenizer if you get string inputs.txt = tokenizer(txt)txt = self.seq_embedding(txt)# Look at the imagefor dec_layer in self.decoder_layers:txt = dec_layer(inputs=(image, txt))txt = self.output_layer(txt)return txt

生成图像的描述：

• 用[START]初始化input_tokens
• 输入(image, input_tokens)到模型,循环生成token：
  • 模型输出下一个token的logits
  • 根据logits采样到下一个token
  • 添加到input_tokens，得到新的input_tokens
  • 如果token==[END]推出循环
• temperature : 控制token生成的采样方式

@Captioner.add_method
def simple_gen(self, image, temperature=1):initial = self.word_to_index([['[START]']]) # (batch, sequence)img_features = self.feature_extractor(image[tf.newaxis, ...])tokens = initial # (batch, sequence)for n in range(50):preds = self((img_features, tokens)).numpy()  # (batch, sequence, vocab)preds = preds[:,-1, :]  #(batch, vocab)if temperature==0:next_token = tf.argmax(preds, axis=-1)[:, tf.newaxis]  # (batch, 1)else:next_token = tf.random.categorical(preds/temperature, num_samples=1)  # (batch, 1)tokens = tf.concat([tokens, next_token], axis=1) # (batch, sequence) if next_token[0] == self.word_to_index('[END]'):breakwords = self.index_to_word(tokens[0, 1:-1])result = tf.strings.reduce_join(words, axis=-1, separator=' ')return result.numpy().decode()

model = Captioner(tokenizer, feature_extractor=mobilenet, output_layer=output_layer,units=256, dropout_rate=0.5, num_layers=2, num_heads=2)

训练模型

损失和目标函数

def masked_loss(labels, preds):  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels, preds)mask = (labels != 0) & (loss < 1e8) mask = tf.cast(mask, loss.dtype)loss = loss*maskloss = tf.reduce_sum(loss)/tf.reduce_sum(mask)return lossdef masked_acc(labels, preds):mask = tf.cast(labels!=0, tf.float32)preds = tf.argmax(preds, axis=-1)labels = tf.cast(labels, tf.int64)match = tf.cast(preds == labels, mask.dtype)acc = tf.reduce_sum(match*mask)/tf.reduce_sum(mask)return acc

回调函数

查看模型训练期间的训练反馈，随机选择一张测试图片，在每一轮训练结束之后，输出模型生成的文本描述信息：

class GenerateText(tf.keras.callbacks.Callback):def __init__(self, image):self.image = image/255.def on_epoch_end(self, epochs=None, logs=None):print()for t in (0.0, 0.5, 1.0):result = self.model.simple_gen(self.image, temperature=t)print(result)print()

回调测试用例：

回调函数测试：

g = GenerateText(test_img)
g.model = model
g.on_epoch_end(0)

的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的 的
的
人前 的 有 两个 一个 旁站 男士 旁边 田野 肩膀 球衣 红色 有 里 有 的

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),loss=masked_loss,metrics=[masked_acc])

callbacks = [GenerateText(test_img),tf.keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True)]

history = model.fit(train_ds.repeat(),steps_per_epoch=5000,validation_data=test_ds.repeat(),validation_steps=100,epochs=100,callbacks=callbacks)

Epoch 1/100
5000/5000 [==============================] - ETA: 0s - loss: 2.0574 - masked_acc: 0.5664
篮球场 上 有 两个 穿着 运动服 的 男人 在 打篮球
篮球场 上 有 两个 穿着 球衣 的 男人 在 打篮球
篮球场 上 有 有 三个 穿着 球服 的 男人 在 打篮球5000/5000 [==============================] - 445s 86ms/step - loss: 2.0574 - masked_acc: 0.5664 - val_loss: 1.6012 - val_masked_acc: 0.6220
Epoch 2/100
5000/5000 [==============================] - ETA: 0s - loss: 1.5945 - masked_acc: 0.6207
两个 穿着 球衣 的 男人 在 整洁 的 运动场 上 打篮球
平坦 的 球场上 有 两位 穿着 球服 的 男士 在 争抢 篮球
球场上 有 两个 衣着 运动服 的 男人 在 打篮球5000/5000 [==============================] - 419s 84ms/step - loss: 1.5945 - masked_acc: 0.6207 - val_loss: 1.4775 - val_masked_acc: 0.6355
Epoch 3/100
5000/5000 [==============================] - ETA: 0s - loss: 1.5014 - masked_acc: 0.6312
两个 穿着 球衣 的 运动员 在 球场上 打篮球
两个 身穿 球衣 的 运动员 在 球场上 打篮球
球场上 有 两个 穿着 球衣 的 运动员 在 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.5014 - masked_acc: 0.6312 - val_loss: 1.4660 - val_masked_acc: 0.6334
Epoch 4/100
5000/5000 [==============================] - ETA: 0s - loss: 1.4544 - masked_acc: 0.6373
篮球场 上 有 两个 穿着 运动服 的 男人 在 打篮球
球场上 有 两个 穿着 运动服 的 男人 在 打篮球
两个 穿着 球衣 的 男人 在 平坦 的 运动场 上 打篮球5000/5000 [==============================] - 418s 84ms/step - loss: 1.4544 - masked_acc: 0.6373 - val_loss: 1.4273 - val_masked_acc: 0.6406
Epoch 5/100
5000/5000 [==============================] - ETA: 0s - loss: 1.4198 - masked_acc: 0.6417
两个 穿着 球衣 的 运动员 在 球场上 打篮球
两个 穿着 不同 颜色 球衣 的 男人 在 篮球场 上 争抢 篮球
篮球场 上 一个 穿着 白色 上衣 的 运动员 在 打篮球5000/5000 [==============================] - 416s 83ms/step - loss: 1.4198 - masked_acc: 0.6417 - val_loss: 1.3511 - val_masked_acc: 0.6541
Epoch 6/100
5000/5000 [==============================] - ETA: 0s - loss: 1.4007 - masked_acc: 0.6438
两个 穿着 球衣 的 运动员 在 球场上 打篮球
两个 穿着 球衣 的 运动员 在 球场上 打篮球
两个 穿着 不同 球衣 的 男人 在 球场上 争抢 篮球5000/5000 [==============================] - 414s 83ms/step - loss: 1.4007 - masked_acc: 0.6438 - val_loss: 1.3481 - val_masked_acc: 0.6513
Epoch 7/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3799 - masked_acc: 0.6470
篮球场 上 有 两个 穿着 不同 球衣 的 男人 在 打篮球
篮球场 上 有 两个 穿着 球衣 的 男人 在 打篮球
两个 穿着 不同 颜色 球衣 的 男人 在 球场上 打篮球5000/5000 [==============================] - 418s 84ms/step - loss: 1.3799 - masked_acc: 0.6470 - val_loss: 1.3453 - val_masked_acc: 0.6508
Epoch 8/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3550 - masked_acc: 0.6502
篮球场 上 有 两个 穿着 运动服 的 男人 在 打篮球
两个 穿着 球衣 的 男人 在 球场上 打篮球
篮球场 上 一个 人 前面 有 三个 穿 运动衣 的 男人 在 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.3550 - masked_acc: 0.6502 - val_loss: 1.3299 - val_masked_acc: 0.6581
Epoch 9/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3423 - masked_acc: 0.6523
篮球场 上 有 两个 穿着 不同 球衣 的 男人 在 打篮球
两个 人 前面 有 两个 穿着 不同 颜色 球衣 的 男人 在 球场上 抢 篮球
两个 穿着 球衣 的 男人 在 球场上 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.3423 - masked_acc: 0.6523 - val_loss: 1.3069 - val_masked_acc: 0.6592
Epoch 10/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3362 - masked_acc: 0.6523
两个 穿着 球衣 的 男人 在 球场上 打篮球
两个 穿着 球衣 的 男人 在 运动场 上 打篮球
球场上 有 两个 穿着 球衣 的 男人 在 打篮球5000/5000 [==============================] - 413s 83ms/step - loss: 1.3362 - masked_acc: 0.6523 - val_loss: 1.2940 - val_masked_acc: 0.6625
Epoch 11/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3222 - masked_acc: 0.6550
两个 穿着 球衣 的 男人 在 球场上 打篮球
两个 穿着 球服 的 男人 在 球场上 打篮球
三个 穿着 球衣 的 男人 在 球场上 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.3222 - masked_acc: 0.6550 - val_loss: 1.2856 - val_masked_acc: 0.6634
Epoch 12/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3148 - masked_acc: 0.6567
两个 穿着 球衣 的 男人 在 球场上 打篮球
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
球场上 有 两个 穿着 球衣 的 男人 在 争抢 篮球5000/5000 [==============================] - 413s 83ms/step - loss: 1.3148 - masked_acc: 0.6567 - val_loss: 1.2696 - val_masked_acc: 0.6615
Epoch 13/100
5000/5000 [==============================] - ETA: 0s - loss: 1.3081 - masked_acc: 0.6576
两个 穿着 不同 颜色 球衣 的 男人 在 球场上 打篮球
两个 穿着 运动服 的 男人 在 球场上 打篮球
篮球场 上 一个 穿着 球服 的 男人 前面 有 两个 穿着 不同 颜色 球衣 的 男人 在 抢球5000/5000 [==============================] - 414s 83ms/step - loss: 1.3081 - masked_acc: 0.6576 - val_loss: 1.2535 - val_masked_acc: 0.6693
Epoch 14/100
5000/5000 [==============================] - ETA: 0s - loss: 1.2994 - masked_acc: 0.6584
两个 穿着 球衣 的 男人 在 球场上 打篮球
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
两个 穿着 运动服 的 男人 在 篮球场 上 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.2994 - masked_acc: 0.6584 - val_loss: 1.3168 - val_masked_acc: 0.6540
Epoch 15/100
5000/5000 [==============================] - ETA: 0s - loss: 1.2869 - masked_acc: 0.6599
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
篮球场 上 有 两个 穿着 运动服 的 男人 在 打篮球
两个 右手 叉腰 的 男人 在 球场上 打篮球5000/5000 [==============================] - 416s 83ms/step - loss: 1.2869 - masked_acc: 0.6599 - val_loss: 1.3368 - val_masked_acc: 0.6497
Epoch 16/100
5000/5000 [==============================] - ETA: 0s - loss: 1.2820 - masked_acc: 0.6608
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
篮球场 上 有 两个 穿着 运动服 的 男人 在 打球
三个 穿着 球服 的 女人 在 球场上 打篮球5000/5000 [==============================] - 411s 82ms/step - loss: 1.2820 - masked_acc: 0.6608 - val_loss: 1.2857 - val_masked_acc: 0.6609
Epoch 17/100
5000/5000 [==============================] - ETA: 0s - loss: 1.2826 - masked_acc: 0.6608
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
两个 穿着 运动服 的 男人 在 运动场 上 打篮球
宽敞 的 球场上 有 两个 身穿 运动服 的 男人 在 打篮球5000/5000 [==============================] - 415s 83ms/step - loss: 1.2826 - masked_acc: 0.6608 - val_loss: 1.3195 - val_masked_acc: 0.6561
Epoch 18/100
5000/5000 [==============================] - ETA: 0s - loss: 1.2720 - masked_acc: 0.6628
两个 穿着 球衣 的 男人 在 球场上 打篮球
两个 穿着 不同 球衣 的 男人 在 运动场 上 争抢 篮球
两个 穿着 运动衣 的 男人 在 运动场 上 抢 篮球5000/5000 [==============================] - 417s 83ms/step - loss: 1.2720 - masked_acc: 0.6628 - val_loss: 1.2930 - val_masked_acc: 0.6601

可视化注意力

# 设置中文字体
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='~/CNfont/chinese_pop.ttf', size=15)

def plot_attention_maps(image, str_tokens, attention_map):fig = plt.figure(figsize=(16, 9))len_result = len(str_tokens)titles = []for i in range(len_result):map = attention_map[i]grid_size = max(int(np.ceil(len_result/3)), 3)ax = fig.add_subplot(3, grid_size, i+1)titles.append(ax.set_title(str_tokens[i], fontproperties=font))img = ax.imshow(image)ax.imshow(map, cmap='gray', alpha=0.5, extent=img.get_extent(),clim=[0.0, np.max(map)])plt.tight_layout()

随机测试：

testfnames = os.listdir('./caption_validation/caption_validation_images/') 
test_img = load_test_image('./caption_validation/caption_validation_images/'+np.random.choice(testfnames))

result = model.simple_gen(test_img/255., temperature=0.0)
result

'一个 双手 拿 着 球杆 的 男人 站 在 高尔夫球场 上'

str_tokens = result.split()
str_tokens.append('[END]')

注意力矩阵：

attn_maps = [layer.last_attention_scores for layer in model.decoder_layers]
[map.shape for map in attn_maps]

[TensorShape([1, 2, 12, 49]), TensorShape([1, 2, 12, 49])]

# 在batch，head 维度上计算注意力均值
attention_maps = tf.concat(attn_maps, axis=0)
attention_maps = einops.reduce(attention_maps,'batch heads sequence (height width) -> sequence height width',height=7, width=7,reduction='mean')attention_maps.shape

TensorShape([12, 7, 7])

每次词对应的注意力权重：

plot_attention_maps(test_img/255, str_tokens, attention_maps)