用PyTorch给图片生成标题
本文,我将描述自动图像标题背后的算法,使用深度学习库 - PyTorch来构建体系结构。
问题陈述
首先,我们需要解决的问题:
给定图像,我们想要获得描述图像组成的句子。
现在,我们可以看到我们的机器学习模型应该将一组图像作为输入并输出一组句子作为输出。神经网络是完成此类任务的完美机器学习系列。
数据集
我将使用COCO(Common Objects in Context)机器学习数据集来训练模型。COCO是用于此类任务的常用数据集。每张图片都有5种不同的字幕,张图片的标题与其他图片的标题略有不同(有时差别很大)。以下是COCO数据集中的数据点示例:
由于提供的目标和数据的多样性,COCO是针对其他竞赛(例如语义分段或对象检测)的强大数据集。
机器学习模型定义
我们的输入是图像,输出将是句子。我们可以把句子看作单词的序列。幸运的是,序列模型可以帮助我们处理单词(或字符,或其他序列数据,如时间序列)的序列。
因此,我们可以将数据从图像空间映射到一些隐藏空间,然后将隐藏空间映射到句子空间。
以下是我们将用于构建机器学习模型的架构概述:
标题模型架构概述
如上图所示,我们将把整个机器学习模型分解为编码器和解码器模型,它们通过一个潜在空间向量进行通信。这种结构很容易理解为函数,即通过编码将图像映射到某个难以处理的潜在空间,通过解码将图像的潜在空间表示映射到句子空间。
编码器
我们将使用卷积神经网络对图像进行编码。关键是要理解,我们可以使用在ImageNet数据集上预训练网络来进行迁移学习。考虑到ImageNet和COCO的数据生成分布不同,整个模型的性能可能会低于平均水平。因此,建议在资源非常有限的情况下进行迁移学习。
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
您应该注意的一个细节是编码器网络的输出维度。注意,网络在潜在空间中产生一个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的隐藏状态输出
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。
- 重复此操作,直到获得<end>令牌作为cell的输出。
为了总结上述算法,我们在下一个时间步中将最可能的单词或字符作为输入传递给LSTM Cell并重复该过程。
然而,深度学习实践者提出了一种称为teacher forcer的算法,并且在大多数情况下(在适用的情况下),它有助于循环神经网络的收敛。重要的是要记住,我们可以将整个标题(句子)作为目标,而不仅仅是部分或单个单词。
teacher forcer算法可以总结如下:
- 在time t = 0时,将(句子或单词的开头)标记作为输入馈送到LSMT cell。
- 使用argmax在time t = 0找到最可能的字符(单词)的索引。
- 在time t = 1时,将下一个标记(来自目标的下一个嵌入字)提供给LSMT cell作为输入。
- 重复此操作,直到获得<end>令牌作为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
训练
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(' ' + stats, end="") sys.stdout.flush() # Save the weights. if epoch % save_every == 0: print(" Saving 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))
请注意,我们有两个模型组件(即编码器和解码器),我们通过将编码器的输出(即潜在空间向量)传递给解码器(即循环神经网络)来共同训练它们。
我在NVIDIA GTX 1080Ti上训练模型,batch size为48,3个epochs,大约需要1天。在3个epochs之后,模型的结果已经非常好。
结果
以下是在COCO数据集的验证部分运行模型的一些结果:
以下是一些照片标题示例:
值得一提的是,可以使用Beam Search实现采样步骤,以获得更好的标题多样性。还有一些注意力机制可能有助于形成更好的标题,因为注意机力制对图像的不同部分给予不同程度的注意。