圖像描述技術，是指生成連貫流暢的語句描述圖像內容。對網際網路中圖像信息的檢索、兒童的早期教育、視障人士的生活輔助等方面都有重要的意義。因此，圖像描述領域受到越來越多的關注。
本案例將採用帶有注意力機制的編碼器-解碼器網絡結構，生成語句來描述圖像內容。
目錄1. 數據集簡介
2. 模型介紹
3. 模型結構
3.1 編碼器
3.2 注意力機制
3.3 解碼器
4. 構建模型
5. 圖像描述生成
6. 總結
本案例採用的模型，已經在MSCOCO 2014數據集上進行了預訓練。MSCOCO 2014的訓練集包含82783張圖像，大小為13.5GB，驗證集包含40504張圖像，大小為6.6GB。
這一數據集收集並標註複雜的日常場景圖像，常用於訓練圖像描述的模型。圖像標註出80類物體，如：人、車輛類、動物類等，同時包含對圖像的文本描述。
2 模型介紹帶有注意力機制的編碼器-解碼器結構：編碼器能夠將輸入圖像轉為向量；解碼器能夠將向量轉化為語句，描述輸入圖像。
引入注意力機制，解碼器能夠在逐字生成語句時，把注意力集中於圖像中與當前詞最相關的部分。
下圖展示了該網絡結構的主要思想：
接下來我們分別構建編碼器和解碼器網絡。
3.1 編碼器首先構建編碼器（Encoder），本案例中的編碼器為在MSCOCO 2014數據集上進行預訓練的ResNet-101網絡。
先加載需要使用的庫：
cp models.py ../import models
import torch
from torch import nn
import torchvision
import torch.nn.functional as F
import numpy as np
import json
import os
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import skimage.transform
from skimage import img_as_ubyte
from skimage.transform import resize
from imageio import imread
from PIL import Image
import random
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")編碼器使用預訓練的模型，可以節省訓練時間。預訓練的ResNet-101結構中，最後一層池化層和全連接層用於圖像分類，而編碼過程不需要進行圖像分類，因此我們刪除這兩層。
增加自適應池化層，將特徵向量調整為一致大小，使模型可以接收任意像素大小的輸入圖像。此外，在網絡結構內，添加對ResNet-101網絡2-4層卷積層進行微調的模塊。
# 編碼器
class Encoder(nn.Module):
def __init__(self, encoded_image_size=14): # encoded_image_size: 生成的特徵圖的大小
super(Encoder, self).__init__()
self.enc_image_size = encoded_image_size
# 加載預訓練的ResNet-101模型
resnet = torchvision.models.resnet101(pretrained=True)
# 刪去最後一層池化層和全連接層（該層用於分類）
modules = list(resnet.children())[:-2]
self.resnet = nn.Sequential(*modules)
# 調整特徵向量大小
self.adaptive_pool = nn.AdaptiveAvgPool2d((encoded_image_size, encoded_image_size))
self.fine_tune()
# 前向傳播 傳入圖片，提取特徵
def forward(self, images):
# 輸出特徵大小：(2048, image_size/32, image_size/32)
out = self.resnet(images)
# 轉化為一致大小：(2048, encoded_image_size, encoded_image_size)
out = self.adaptive_pool(out)
# 調換位置 (encoded_image_size, encoded_image_size, 2048)
out = out.permute(0, 2, 3, 1)
# 輸出圖像經過編碼後，得到的向量
return out
# 微調模型
def fine_tune(self, fine_tune=True):
for p in self.resnet.parameters():
p.requires_grad = False
# 對2-4層卷積層進行微調
for c in list(self.resnet.children())[5:]:
for p in c.parameters(): # 微調某層參數
p.requires_grad = fine_tune
3.2 注意力機制在構建解碼器前，需要先定義注意力機制（Attention）。注意力機制返回加權的特徵向量。利用加權的特徵向量，與上一個預測的單詞，結合起來預測下一個單詞，解碼器能夠生成更加準確的語句。下圖展示了該機制的結構：
# 注意力機制
class Attention(nn.Module):
def __init__(self, encoder_dim, decoder_dim, attention_dim):
# encoder_dim: 編碼圖像的特徵維度 ; decoder_dim: decoder維度; attention_dim: 注意力機制層數
super(Attention, self).__init__()
# 兩個注意力機制模塊，分別針對編碼和解碼
self.encoder_att = nn.Linear(encoder_dim, attention_dim) # 全連接層轉換編碼的特徵向量
self.decoder_att = nn.Linear(decoder_dim, attention_dim) # 全連接層轉換解碼器的輸出
self.full_att = nn.Linear(attention_dim, 1)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1) # softmax層計算權重
# 前向傳播
def forward(self, encoder_out, decoder_hidden):
att1 = self.encoder_att(encoder_out) # 將圖像特徵傳入att1
att2 = self.decoder_att(decoder_hidden) # batch_size * attention_dim
att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2) # 每個像素的權重
alpha = self.softmax(att) # softmax計算權重
attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)
# 返回每個時刻加權的圖像特徵向量
return attention_weighted_encoding, alpha
3.3 解碼器最後構建帶有注意力機制的解碼器（Decoder）。本案例採用LSTM作為解碼器，逐字生成語句。因為注意力機制，解碼器在生成不同單詞的時候，會關注圖像不同部分。比如，生成"a man holds a football"中的"football"時，解碼器會關注圖像中的「足球」。
解碼器保留LSTM得到的中間輸出語句，並在生成新的單詞時，先用注意力機制，得到特徵向量的權重。通過之前預測的語句和新的權重，共同預測下一個單詞。
本案例中，帶有注意力機制的LSTM網絡結構包含：注意力機制、嵌入層、dropout層、通過前向傳播解碼。
class DecoderWithAttention(nn.Module):
def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, encoder_dim=2048, dropout=0.5):
# attention_dim: 注意力機制的層數; embed_dim: 嵌入層; decoder_dim: 解碼器維度
# vocab_size: 單詞表 word map; encoder_dim: 編碼器維度
super(DecoderWithAttention, self).__init__()
self.encoder_dim = encoder_dim
self.attention_dim = attention_dim
self.embed_dim = embed_dim
self.decoder_dim = decoder_dim
self.vocab_size = vocab_size
self.dropout = dropout
self.attention = Attention(encoder_dim, decoder_dim, attention_dim) # 注意力機制
self.embedding = nn.Embedding(vocab_size, embed_dim) # 嵌入層 數據降維、轉換為稠密向量
self.dropout = nn.Dropout(p=self.dropout)
self.decode_step = nn.LSTMCell(embed_dim + encoder_dim, decoder_dim, bias=True)
# LSTM細胞
self.init_h = nn.Linear(encoder_dim, decoder_dim) # 初始隱藏狀態
self.init_c = nn.Linear(encoder_dim, decoder_dim) # 初始細胞狀態
self.f_beta = nn.Linear(decoder_dim, encoder_dim)
self.sigmoid = nn.Sigmoid() # sigmoid型激活門
self.fc = nn.Linear(decoder_dim, vocab_size) # 為單詞打分
self.init_weights() # 初始化權重為均勻分布
# 將權重初始化為均勻分布
def init_weights(self):
self.embedding.weight.data.uniform_(-0.1, 0.1)
self.fc.bias.data.fill_(0)
self.fc.weight.data.uniform_(-0.1, 0.1)
# 加載預訓練的嵌入層
def load_pretrained_embeddings(self, embeddings):
self.embedding.weight = nn.Parameter(embeddings)
# 微調嵌入層
def fine_tune_embeddings(self, fine_tune=True):
for p in self.embedding.parameters():
p.requires_grad = fine_tune
# 隱藏層初始狀態
def init_hidden_state(self, encoder_out):
mean_encoder_out = encoder_out.mean(dim=1) # 對第二維求平均 如果已經展平 相當於求每層所有像素的平均值
h = self.init_h(mean_encoder_out) # 大小為 (batch_size, decoder_dim)
c = self.init_c(mean_encoder_out)
return h, c
# 前向傳播
def forward(self, encoder_out, encoded_captions, caption_lengths):
batch_size = encoder_out.size(0)
encoder_dim = encoder_out.size(-1)
vocab_size = self.vocab_size
# 編碼圖像
encoder_out = encoder_out.view(batch_size, -1, encoder_dim) # 輸出向量：(batch_size, num_pixels, encoder_dim)
num_pixels = encoder_out.size(1)
# 對句子長度進行排序
caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)
encoder_out = encoder_out[sort_ind]
encoded_captions = encoded_captions[sort_ind]
# 嵌入層 將每個詞用512維的向量表示
embeddings = self.embedding(encoded_captions) # 嵌入層
h, c = self.init_hidden_state(encoder_out) # LSTM初始狀態
decode_lengths = (caption_lengths - 1).tolist() # 在<end>部分停止解碼
# max(decode_lengths): 最長詞語數量; vocab_size: 單詞表
predictions = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(device)
# num_pixels: 像素總數
alphas = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(device)
# 用之前的單詞和注意力機製得到的權重，解碼新的單詞
for t in range(max(decode_lengths)): # max(decode_lengths): 句子最長長度
batch_size_t = sum([l > t for l in decode_lengths])
attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t],h[:batch_size_t]) # 每次循環都重新生成注意力權重
gate = self.sigmoid(self.f_beta(h[:batch_size_t])) # (batch_size_t, encoder_dim)
attention_weighted_encoding = gate * attention_weighted_encoding # 每層的注意力權重
h, c = self.decode_step(torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),
(h[:batch_size_t], c[:batch_size_t])) # (batch_size_t, decoder_dim)
preds = self.fc(self.dropout(h)) # (batch_size_t, vocab_size)
predictions[:batch_size_t, t, :] = preds # 生成預測
alphas[:batch_size_t, t, :] = alpha # 每個像素的重點區域
return predictions, encoded_captions, decode_lengths, alphas, sort_ind
4 構建模型構建好上述編碼器、解碼器、注意力機制後，需要定義實現圖像描述的函數caption_image_beam_search，該函數連接編碼器、注意力機制、解碼器。即將圖像傳入編碼器後，計算得到特徵向量；將特徵向量傳入帶有注意力機制的解碼器，通過LSTM算法得到單詞表內生成語句對應的數字。
在搜索最優生成語句時，我們使用beam search（束搜索）算法，即每次挑選出所有生成語句中條件概率最大的k個，作為該時間步長下的候選輸出序列。始終保持k個候選語句，最後從k個候選中挑出最優的生成語句。
def caption_image_beam_search(encoder, decoder, image_path, word_map, beam_size=3):
k = beam_size # 每個解碼步驟中，選擇概率最大的k個詞
vocab_size = len(word_map) # 單詞表個數
# 讀取圖像
img = imread(image_path)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img = np.concatenate([img, img, img], axis=2)
img = img_as_ubyte(resize(img, (256, 256)))
img = img.transpose(2, 0, 1)
img = img / 255.
img = torch.FloatTensor(img).to(device)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([normalize])
image = transform(img) # (3, 256, 256)
# 編碼器
image = image.unsqueeze(0) # (1, 3, 256, 256)
encoder_out = encoder(image) # (1, enc_image_size, enc_image_size, encoder_dim)
enc_image_size = encoder_out.size(1)
encoder_dim = encoder_out.size(3)
# 編碼得到的特徵向量
encoder_out = encoder_out.view(1, -1, encoder_dim) # (1, num_pixels, encoder_dim)
num_pixels = encoder_out.size(1)
# batch size = k
encoder_out = encoder_out.expand(k, num_pixels, encoder_dim) # (k, num_pixels, encoder_dim)
# 每次儲存前k個單詞，初始狀態為<start>
k_prev_words = torch.LongTensor([[word_map['<start>']]] * k).to(device) # (k, 1)
seqs = k_prev_words # (k, 1)
top_k_scores = torch.zeros(k, 1).to(device) # (k, 1)
# 初始權重
seqs_alpha = torch.ones(k, 1, enc_image_size, enc_image_size).to(device)
complete_seqs = list()
complete_seqs_alpha = list()
complete_seqs_scores = list()
# 解碼
step = 1
h, c = decoder.init_hidden_state(encoder_out)
# s <= k, 到達<end>時停止解碼
while True:
embeddings = decoder.embedding(k_prev_words).squeeze(1) # 維度: (s, embed_dim)
awe, alpha = decoder.attention(encoder_out, h) # (s, encoder_dim), (s, num_pixels)
alpha = alpha.view(-1, enc_image_size, enc_image_size) # (s, enc_image_size, enc_image_size)
gate = decoder.sigmoid(decoder.f_beta(h)) # gating scalar, (s, encoder_dim)
awe = gate * awe
h, c = decoder.decode_step(torch.cat([embeddings, awe], dim=1), (h, c)) # (s, decoder_dim)
scores = decoder.fc(h) # (s, vocab_size)
scores = F.log_softmax(scores, dim=1)
scores = top_k_scores.expand_as(scores) + scores # (s, vocab_size)
# 初始狀態下，分數相同
if step == 1:
top_k_scores, top_k_words = scores[0].topk(k, 0, True, True) # (s)
else:
# 找到最高得分的單詞
top_k_scores, top_k_words = scores.view(-1).topk(k, 0, True, True) # (s)
prev_word_inds = top_k_words // vocab_size # (s)
next_word_inds = top_k_words % vocab_size # (s)
# 輸出新單詞
seqs = torch.cat([seqs[prev_word_inds], next_word_inds.unsqueeze(1)], dim=1) # (s, step+1)
seqs_alpha = torch.cat([seqs_alpha[prev_word_inds], alpha[prev_word_inds].unsqueeze(1)],
dim=1) # (s, step+1, enc_image_size, enc_image_size)
# 判斷是否可以停止解碼
incomplete_inds = [ind for ind, next_word in enumerate(next_word_inds) if
next_word != word_map['<end>']]
complete_inds = list(set(range(len(next_word_inds))) - set(incomplete_inds))
# 完整語句
if len(complete_inds) > 0:
complete_seqs.extend(seqs[complete_inds].tolist())
complete_seqs_alpha.extend(seqs_alpha[complete_inds].tolist())
complete_seqs_scores.extend(top_k_scores[complete_inds])
k -= len(complete_inds)
# 處理不完整的語句
if k == 0:
break
seqs = seqs[incomplete_inds]
seqs_alpha = seqs_alpha[incomplete_inds]
h = h[prev_word_inds[incomplete_inds]]
c = c[prev_word_inds[incomplete_inds]]
encoder_out = encoder_out[prev_word_inds[incomplete_inds]]
top_k_scores = top_k_scores[incomplete_inds].unsqueeze(1)
k_prev_words = next_word_inds[incomplete_inds].unsqueeze(1)
# 步驟過長則終止
if step > 50:
break
step += 1
i = complete_seqs_scores.index(max(complete_seqs_scores))
seq = complete_seqs[i]
alphas = complete_seqs_alpha[i]
return seq, alphas為了更好地展示輸出結果，我們定義visualize_att函數，將圖像注意力集中的位置和對應的輸出單詞，直觀展示出來。
def visualize_att(image_path, seq, alphas, rev_word_map, smooth=True):
image = Image.open(image_path) # 讀入圖像
image = image.resize([14 * 24, 14 * 24], Image.LANCZOS) # 調整圖像大小
words = [rev_word_map[ind] for ind in seq]
plt.figure(figsize=(18,9))
for t in range(len(words)):
if t > 50:
break
plt.subplot(np.ceil(len(words) / 5.), 5, t + 1)
# 將注意力集中的圖像部分高亮
plt.text(0, 1, '%s' % (words[t]), color='black', backgroundcolor='white', fontsize=12)
plt.imshow(image)
current_alpha = alphas[t, :]
if smooth:
alpha = skimage.transform.pyramid_expand(current_alpha.numpy(), upscale=24, sigma=8)
else:
alpha = skimage.transform.resize(current_alpha.numpy(), [14 * 24, 14 * 24])
if t == 0:
plt.imshow(alpha, alpha=0)
else:
plt.imshow(alpha, alpha=0.8)
plt.set_cmap(cm.Greys_r)
plt.axis('off')
plt.show()
5 圖像描述生成現在我們可以通過調用上述定義的函數，對輸入的圖像生成對應的語句。考慮到訓練時間過長，本案例將加載訓練好的編碼器與解碼器參數。
from glob import glob
checkpoint = '/content/BEST_checkpoint_coco_5_cap_per_img_5_min_word_freq.pth.tar' # 預訓練參數
word_map_file = '/content/WORDMAP_coco_5_cap_per_img_5_min_word_freq.json' # 單詞表
beam_size = 5
smooth = True
# 加載模型
checkpoint = torch.load(checkpoint, map_location=device)
decoder = checkpoint['decoder']
decoder = decoder.to(device)
decoder.eval()
encoder = checkpoint['encoder']
encoder = encoder.to(device)
encoder.eval()加載單詞表，即每個數字對應的單詞。
# 加載單詞表
with open(word_map_file, 'r') as j:
word_map = json.load(j)
# 轉換格式
rev_word_map = {v: k for k, v in word_map.items()}
# 展示單詞表
for i in range(5):
t = random.choice(list(rev_word_map))
print(t,":",rev_word_map[t])333 : chairs
505 : plant
1575 : snowboarder
1859 : indian
3049 : build接下來，我們可以向模型輸入一張圖像進行描述，我們使用的數據如下圖所示：
Image.open("/content/surf.png") # 打開圖像my_files = np.array(glob("/content/surf.png"))
for img in my_files:
# 生成圖像描述
seq, alphas = caption_image_beam_search(encoder, decoder, img, word_map, beam_size)
alphas = torch.FloatTensor(alphas)
for i in seq:
print(rev_word_map[i],end=' ')
# 展示圖像描述結果
visualize_att(img, seq, alphas, rev_word_map, smooth)<start> a couple of people with surfboards in the water <end>生成的語句為："a couple of people with surfboards in the water"，即「海裡有兩個人拿著衝浪板」，語句通順，且準確地描述出了圖像的內容，包含了「衝浪板」、「兩個人」、「海」這些關鍵詞語。
通過可視化注意力機制，我們看到在生成"people"這一單詞時，模型關注圖像中的「兩個人」；生成"surfboards"時，則關注人們手中的「衝浪板」。
6 總結本案例構建了帶有注意力機制的ResNet-101和LSTM的模型，實現圖像描述。其中ResNet-101模型已經在MSCOCO 2014上進行了預訓練，節省了模型訓練的時間。在尋找最優語句時，使用beam search算法，實現搜索最優語句的任務。
觀察生成結果，可知模型的訓練結果較好，生成的語句能夠準確地描述圖像內容。