用PyTorch給圖片生成標題

本文,我將描述自動圖像標題背後的算法,使用深度學習庫 - PyTorch來構建體系結構。

問題陳述

首先,我們需要解決的問題:

給定圖像,我們想要獲得描述圖像組成的句子。

現在,我們可以看到我們的機器學習模型應該將一組圖像作為輸入並輸出一組句子作為輸出。神經網絡是完成此類任務的完美機器學習系列。

數據集

我將使用COCO(Common Objects in Context)機器學習數據集來訓練模型。COCO是用於此類任務的常用數據集。每張圖片都有5種不同的字幕,張圖片的標題與其他圖片的標題略有不同(有時差別很大)。以下是COCO數據集中的數據點示例:

用PyTorch給圖片生成標題

由於提供的目標和數據的多樣性,COCO是針對其他競賽(例如語義分段或對象檢測)的強大數據集。

機器學習模型定義

我們的輸入是圖像,輸出將是句子。我們可以把句子看作單詞的序列。幸運的是,序列模型可以幫助我們處理單詞(或字符,或其他序列數據,如時間序列)的序列。

因此,我們可以將數據從圖像空間映射到一些隱藏空間,然後將隱藏空間映射到句子空間。

以下是我們將用於構建機器學習模型的架構概述:

用PyTorch給圖片生成標題

標題模型架構概述

如上圖所示,我們將把整個機器學習模型分解為編碼器和解碼器模型,它們通過一個潛在空間向量進行通信。這種結構很容易理解為函數,即通過編碼將圖像映射到某個難以處理的潛在空間,通過解碼將圖像的潛在空間表示映射到句子空間。

編碼器

我們將使用卷積神經網絡對圖像進行編碼。關鍵是要理解,我們可以使用在ImageNet數據集上預訓練網絡來進行遷移學習。考慮到ImageNet和COCO的數據生成分佈不同,整個模型的性能可能會低於平均水平。因此,建議在資源非常有限的情況下進行遷移學習。

用PyTorch給圖片生成標題

Dense block: arXiv:1608.06993v5

對於這個特定模型,我在編碼器中使用了DenseNet121架構。它是一種相對輕量級且性能良好的計算機視覺應用架構; 您也可以使用任何其他卷積架構將圖像映射到潛在空間。我使用的dense網絡沒有預訓練以避免數據生成分佈之間的轉換。

我們可以通過使用torchvision的models 包,輕鬆地導入PyTorch中的模型。即使導入的dense network沒有經過預先訓練,它仍然帶有ImageNet數據集的分類器,該機器學習數據集有1000個類。幸運的是,在PyTorch中很容易替換分類器。我用參數ReLU激活函數的雙層感知器代替了網絡的分類器,並使用dropout減少過度擬合。你可以在PyTorch中找到編碼器的實現,Python代碼如下:

class EncoderCNN(nn.Module):
def __init__(self, embed_size = 1024):
super(EncoderCNN, self).__init__()

# get the pretrained densenet model
self.densenet = models.densenet121(pretrained=True)

# replace the classifier with a fully connected embedding layer
self.densenet.classifier = nn.Linear(in_features=1024, out_features=1024)

# add another fully connected layer
self.embed = nn.Linear(in_features=1024, out_features=embed_size)

# dropout layer
self.dropout = nn.Dropout(p=0.5)

# activation layers
self.prelu = nn.PReLU()

def forward(self, images):

# get the embeddings from the densenet
densenet_outputs = self.dropout(self.prelu(self.densenet(images)))

# pass through the fully connected
embeddings = self.embed(densenet_outputs)

return embeddings
用PyTorch給圖片生成標題

您應該注意的一個細節是編碼器網絡的輸出維度。注意,網絡在潛在空間中產生一個1024維的向量,我們將其作為LSTM模型的第一個輸入(at time t=0)。

解碼器

LSTM cell

到目前為止,一切都很簡單:我們有一個圖像,我們通過一個稍微修改過的緊密相連的神經網絡,得到一個1024維的輸出向量。解碼器是體系結構的一部分。

正如我在架構概述中所展示的那樣,解碼器由一個循環神經網絡組成。我們可以使用GRU或LSTM單元,這裡使用了後者。

LSTM cell具有長期和短期內存(duh)。

  • LSTM cell具有數據點的輸入(at time t=n)
  • LSTM cell具有用於cell狀態的輸入(先前cell狀態)
  • LSTM cell具有隱藏狀態的輸入(先前的隱藏狀態)
  • LSTM cell具有cell狀態的輸出(當前cell狀態)
  • LSTM cell具有隱藏狀態的輸出(當前隱藏狀態)
  • LSTM cell的數據輸出是LSTM cell的隱藏狀態輸出
用PyTorch給圖片生成標題

LSTM cell

理解實現更重要的是,對於序列中的每個步驟,我們使用完全相同的LSTM(或GRU)cell,因此優化cell的目標是找到正確的權重集以適應整個單詞詞典(char-to-char模型中的字符)。這意味著對於我們句子中的每個單詞(這是一個序列),我們將把這個單詞作為輸入提供並獲得一些輸出,這通常是整個單詞詞典的概率分佈。通過這種方式,我們可以獲得模型認為最擬合前一個單詞的單詞。

詞典(Dictionary)/詞彙(Vocabulary

序列模型(實際上通常是模型)不理解符號語言,即圖像必須表示為實數的張量,以便模型能夠處理它們,因為神經網絡是在其間具有非線性的多個並行(向量化)計算。將圖像轉換為模型理解的語言非常簡單,最常見的方法是採用實數表示的每個像素的強度。幸運的是,有一種方法可以將單詞轉換為這種語言。

現在,我們知道我們的數據生成分佈可以作為目標句的一部分生成的有意義的單詞的數量是有限的。所以我們要做的是將訓練數據集中所有標題中出現的每一個單詞進行枚舉,以獲得單詞和整數之間的映射。我們已經完成了一半,可以開始使用解碼器中的單詞了。現在我們可以構建一個1到k的映射(通常使用單層感知器),將單詞的整數表示形式映射到一個k維空間,我們可以將這個空間用作LSTM cell的輸入。

我們現在可以看到,每個單詞都將嵌入到更高維度的實數空間中,我們可以使用它來處理循環神經網絡。嵌入在其他自然語言處理應用中也是有用的,因為它們允許從業者一旦被映射到2維空間(通常使用t - sne算法)來檢查單詞或字符集合。

Teacher Forcer

以下是使用循環網絡的常見方案:

  • 在time t = 0時,將(句子或單詞的開頭)標記作為輸入饋送到LSMT cell。
  • 在time t = 0時,獲得詞彙量大小的向量作為LSTM的輸出。
  • 使用argmax在time t=0時查找最可能字符(單詞)的索引。
  • 嵌入最可能的字符(單詞)。
  • 在t =1時將結果嵌入作為輸入傳遞給LSTM cell。
  • 重複此操作,直到獲得令牌作為cell的輸出。

為了總結上述算法,我們在下一個時間步中將最可能的單詞或字符作為輸入傳遞給LSTM Cell並重復該過程。

然而,深度學習實踐者提出了一種稱為teacher forcer的算法,並且在大多數情況下(在適用的情況下),它有助於循環神經網絡的收斂。重要的是要記住,我們可以將整個標題(句子)作為目標,而不僅僅是部分或單個單詞。

teacher forcer算法可以總結如下:

  • 在time t = 0時,將(句子或單詞的開頭)標記作為輸入饋送到LSMT cell。
  • 使用argmax在time t = 0找到最可能的字符(單詞)的索引。
  • 在time t = 1時,將下一個標記(來自目標的下一個嵌入字)提供給LSMT cell作為輸入。
  • 重複此操作,直到獲得令牌作為cell的輸出。

請注意,我們不再提供最後一個最可能的單詞,我們提供已經可用的下一個單詞嵌入。

解碼器實現

首先,我們使用單層LSTM將潛在空間向量映射到單詞空間。

其次,正如我前面提到的,LSTM單元的輸出是隱藏狀態向量(在LSTM cell圖中以紫色顯示)。因此,我們需要從隱藏狀態空間到詞彙(字典)空間的某種映射。我們可以通過在隱藏狀態空間和詞彙空間之間使用全連接層來實現這一點。

如果你對循環神經網絡有一定的經驗,Forward pass是很簡單的。

這裡的關鍵思想是在time t = 0時將表示圖像的潛在空間向量作為LSTM單元的輸入。從time t = 1開始,我們可以開始將我們的嵌入式目標句子作為teacher forcer算法的一部分提供給LSTM單元。

class DecoderRNN(nn.Module):
def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1):
super(DecoderRNN, self).__init__()

# define the properties
self.embed_size = embed_size
self.hidden_size = hidden_size
self.vocab_size = vocab_size

# lstm cell
self.lstm_cell = nn.LSTMCell(input_size=embed_size, hidden_size=hidden_size)

# output fully connected layer
self.fc_out = nn.Linear(in_features=self.hidden_size, out_features=self.vocab_size)

# embedding layer
self.embed = nn.Embedding(num_embeddings=self.vocab_size, embedding_dim=self.embed_size)

# activations
self.softmax = nn.Softmax(dim=1)

def forward(self, features, captions):

# batch size
batch_size = features.size(0)

# init the hidden and cell states to zeros
hidden_state = torch.zeros((batch_size, self.hidden_size)).cuda()
cell_state = torch.zeros((batch_size, self.hidden_size)).cuda()

# define the output tensor placeholder
outputs = torch.empty((batch_size, captions.size(1), self.vocab_size)).cuda()
# embed the captions
captions_embed = self.embed(captions)

# pass the caption word by word
for t in range(captions.size(1)):

# for the first time step the input is the feature vector
if t == 0:
hidden_state, cell_state = self.lstm_cell(features, (hidden_state, cell_state))

# for the 2nd+ time step, using teacher forcer
else:
hidden_state, cell_state = self.lstm_cell(captions_embed[:, t, :], (hidden_state, cell_state))

# output of the attention mechanism
out = self.fc_out(hidden_state)

# build the output tensor
outputs[:, t, :] = out

return outputs
用PyTorch給圖片生成標題

訓練

Python代碼如下:

# get the losses for vizualization
losses = list()
val_losses = list()
for epoch in range(1, 10+1):

for i_step in range(1, total_step+1):

# zero the gradients
decoder.zero_grad()
encoder.zero_grad()

# set decoder and encoder into train mode
encoder.train()
decoder.train()

# Randomly sample a caption length, and sample indices with that length.
indices = train_data_loader.dataset.get_train_indices()

# Create and assign a batch sampler to retrieve a batch with the sampled indices.
new_sampler = data.sampler.SubsetRandomSampler(indices=indices)
train_data_loader.batch_sampler.sampler = new_sampler

# Obtain the batch.
images, captions = next(iter(train_data_loader))

# make the captions for targets and teacher forcer
captions_target = captions[:, 1:].to(device)
captions_train = captions[:, :captions.shape[1]-1].to(device)
# Move batch of images and captions to GPU if CUDA is available.
images = images.to(device)

# Pass the inputs through the CNN-RNN model.
features = encoder(images)
outputs = decoder(features, captions_train)

# Calculate the batch loss
loss = criterion(outputs.view(-1, vocab_size), captions_target.contiguous().view(-1))

# Backward pass
loss.backward()

# Update the parameters in the optimizer
optimizer.step()

# - - - Validate - - -
# turn the evaluation mode on
with torch.no_grad():

# set the evaluation mode

encoder.eval()
decoder.eval()
# get the validation images and captions
val_images, val_captions = next(iter(val_data_loader))
# define the captions
captions_target = val_captions[:, 1:].to(device)
captions_train = val_captions[:, :val_captions.shape[1]-1].to(device)
# Move batch of images and captions to GPU if CUDA is available.
val_images = val_images.to(device)
# Pass the inputs through the CNN-RNN model.
features = encoder(val_images)
outputs = decoder(features, captions_train)
# Calculate the batch loss.
val_loss = criterion(outputs.view(-1, vocab_size), captions_target.contiguous().view(-1))

# append the validation loss and training loss
val_losses.append(val_loss.item())
losses.append(loss.item())

# save the losses
np.save('losses', np.array(losses))
np.save('val_losses', np.array(val_losses))

# Get training statistics.
stats = 'Epoch [%d/%d], Step [%d/%d], Loss: %.4f, Val Loss: %.4f' % (epoch, num_epochs, i_step, total_step, loss.item(), val_loss.item())

# Print training statistics (on same line).
print('\\r' + stats, end="")
sys.stdout.flush()

# Save the weights.
if epoch % save_every == 0:
print("\\nSaving the model")
torch.save(decoder.state_dict(), os.path.join('./models', 'decoder-%d.pth' % epoch))
torch.save(encoder.state_dict(), os.path.join('./models', 'encoder-%d.pth' % epoch))
用PyTorch給圖片生成標題

用PyTorch給圖片生成標題

請注意,我們有兩個模型組件(即編碼器和解碼器),我們通過將編碼器的輸出(即潛在空間向量)傳遞給解碼器(即循環神經網絡)來共同訓練它們。

我在NVIDIA GTX 1080Ti上訓練模型,batch size為48,3個epochs,大約需要1天。在3個epochs之後,模型的結果已經非常好。

結果

以下是在COCO數據集的驗證部分運行模型的一些結果:

用PyTorch給圖片生成標題

以下是一些照片標題示例:

用PyTorch給圖片生成標題

值得一提的是,可以使用Beam Search實現採樣步驟,以獲得更好的標題多樣性。還有一些注意力機制可能有助於形成更好的標題,因為注意機力制對圖像的不同部分給予不同程度的注意。


分享到:


相關文章: