图像描述(image caption)模型简单demo(源码理解原理)

文章目录

前言
一、推理方式
- 1、逐字推理（word-by-word inference）
- 2、完全推理（full-sentence inference）
二、数据定义
二、模型搭建
三、模型训练
四、模型推理
五、所有Demo源码
总结

前言

图像描述模型简单理解是通过图像给出一段描述内容，实际也类似一种生成模式，和我一篇博客讲的对话模型相似。然而，网络很少从代码层次去说明图像描述具体做法与细节。基于此，我原创一个简单模型，旨在帮助理解原理，但不具备很好效果能力。

一、推理方式

逐字推理（word-by-word inference）和完全推理（full-sentence inference）是在自然语言处理中常见的两种生成文本的方法。

1、逐字推理（word-by-word inference）

在逐字推理中，生成文本的过程是逐字逐字地进行的。模型在每个时间步骤根据前一个时间步骤生成的单词来预测下一个单词。这种方法通常使用循环神经网络（RNN）或者长短时记忆网络（LSTM）等模型来实现。模型在每个时间步骤都会输出一个单词，直到生成完整的句子或段落。

2、完全推理（full-sentence inference）

在完全推理中，生成文本的过程是一次性完成的。模型一次性接收所有输入，并在单个步骤中生成完整的句子或段落。这种方法通常使用注意力机制（attention mechanism）等技术来帮助模型在处理整个输入时关注到不同部分的信息。完全推理通常比逐字推理更高效，因为它可以并行地处理整个输入。

总的来说，逐字推理逐步生成文本，而完全推理一次性给出完整的文本。选择使用哪种方法取决于任务的需求和模型的设计。

二、数据定义

涉及到文本相关内容，文本数据最终都是转为对应字典索引代表其文本内容，输入模型加工，实现文本相关任务，图像描述模型也不列外。因此，我们需要构建一个字典映射(可参考：点击这里)与文本数据，并按照字典映射转换为对应索引id，其代码如下：

# 我这里假设构造一个图像特征是一个且对应描述也是一个
vocab_size = 11  # 字典大小，也是后面概率预测数，一般是32000
vocab = {0:"pad",1: "start", 2: "是", 3: '我', 4: "描述", 5: "图像", 6: "模型", 7: "s", 8: "1", 9: "等", 10: "end"}  # 字典大小为9
images_rand = torch.randn(1, 2048)  # 随机生成一张图像特征
caption = torch.tensor([0, 2, 1, 4, 3, 5, 9]).reshape(1,-1)  # 我是图像描述模型
start_vocab = 1
end_vocab = 10
max_length = 20  # 图像长度

当然，我也定义了相关词汇变量。同时，我们也给出类似dataset加工数据模块，其代码如下：

# 定义随机生成数据的 Dataset,我将其固定位一个图像
class RandomDataset(Dataset):
    def __init__(self, num_samples, max_length):
        # self.images = torch.randn(num_samples, 2048)  # 随机生成图像特征
        self.images = images_rand  # 随机生成图像特征

        self.captions = F.pad(caption, (0, max_length-max(caption.shape)), value=0)

        self.num_samples = num_samples
    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # 因为我只想使用一张图与一个描述，因此idx只起到占位
        return self.images[0], self.captions[0]

二、模型搭建

这里，也是最重要内容，如何搭建图像描述模型，我是使用transformer结构搭建，创建一个简单的图像描述模型，模型共包含2个部分内容，其一是图形特征提取，其二是图形文本生成内容。

1、图像特征提取

为了更加简单，我使用一个线性方式来做为图像特征提取，一般图像输入为[b,c,h,w]，而我为了更简单，直接创建[b,c]维度，表明图像已被处理了。其代码如下：


# 定义图像编码器
class ImageEncoder(nn.Module):
    def __init__(self, embed_size):
        super(ImageEncoder, self).__init__()
        self.linear = nn.Linear(2048, embed_size)  # 假设输入特征维度为2048

    def forward(self, images):
        features = self.linear(images)
        return features

2、图像描述文本解码

这个很重要，实际有点类似transformer的解码结构，只不过文本id特征作为了query而图像特征作为k与v的方式，其代码如下：

# 定义 Transformer 解码器
class TransformerDecoder(nn.Module):
    def __init__(self, embed_size, heads, num_layers, vocab_size):
        super(TransformerDecoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.transformer_layer = nn.TransformerDecoderLayer(d_model=embed_size, nhead=heads)
        self.transformer = nn.TransformerDecoder(self.transformer_layer, num_layers)
        self.linear = nn.Linear(embed_size, vocab_size)

    def forward(self, features, captions):
        embeddings = self.embedding(captions)
        embeddings = embeddings #.permute(1, 0, 2)  # 调整维度顺序
        features = features.unsqueeze(1).repeat(1, captions.shape[1], 1)  # 重复图像特征以匹配序列长度
        output = self.transformer(embeddings, features)
        output = self.linear(output)
        return output

3、图像编码与文本解码结合模型

最终图像编码模型与文本解码模型整合如下代码：

# 定义整合模型
class ImageCaptioningModel(nn.Module):
    def __init__(self, embed_size, heads, num_layers, vocab_size):
        super(ImageCaptioningModel, self).__init__()
        self.image_encoder = ImageEncoder(embed_size)
        self.text_decoder = TransformerDecoder(embed_size, heads, num_layers, vocab_size)

    def forward(self, images, captions):
        features = self.image_encoder(images)
        output = self.text_decoder(features, captions)
        return output

三、模型训练

模型训练莫非就是构建模型、构建数据、构建优化器、给定训练逻辑，我将不在过多解释，直接上代码，如下：

# 创建模型和数据加载器
    model = ImageCaptioningModel(embed_size=256, heads=8, num_layers=2, vocab_size=vocab_size)
    dataset = RandomDataset(num_samples=40, max_length=max_length)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    for epoch in range(60):
        for images, captions in dataloader:
            optimizer.zero_grad()
            features = model.image_encoder(images)
            outputs = model.text_decoder(features, captions)

            loss = criterion(outputs.view(-1, vocab_size), captions.view(-1))  # 计算损失
            loss.backward()
            optimizer.step()

            print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

四、模型推理

最后，让我们使用训练好的模型进行推理，我采用是逐字生成推理方法，正如我上面介绍内容，我也不在解释直接给出代码，如下：

 # 推理
    predicted_caption = inference(images_rand,model)  # images的shape是[1,256]

    print('Predicted Caption:', predicted_caption)
    result = ''.join(vocab[int(index)] for index in predicted_caption)
    print('Predicted Caption:', result)

五、所有Demo源码

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import torch.nn.functional as F


# 定义图像编码器
class ImageEncoder(nn.Module):
    def __init__(self, embed_size):
        super(ImageEncoder, self).__init__()
        self.linear = nn.Linear(2048, embed_size)  # 假设输入特征维度为2048

    def forward(self, images):
        features = self.linear(images)
        return features


# 定义 Transformer 解码器
class TransformerDecoder(nn.Module):
    def __init__(self, embed_size, heads, num_layers, vocab_size):
        super(TransformerDecoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.transformer_layer = nn.TransformerDecoderLayer(d_model=embed_size, nhead=heads)
        self.transformer = nn.TransformerDecoder(self.transformer_layer, num_layers)
        self.linear = nn.Linear(embed_size, vocab_size)

    def forward(self, features, captions):
        embeddings = self.embedding(captions)
        embeddings = embeddings #.permute(1, 0, 2)  # 调整维度顺序
        features = features.unsqueeze(1).repeat(1, captions.shape[1], 1)  # 重复图像特征以匹配序列长度
        output = self.transformer(embeddings, features)
        output = self.linear(output)
        return output


# 定义整合模型
class ImageCaptioningModel(nn.Module):
    def __init__(self, embed_size, heads, num_layers, vocab_size):
        super(ImageCaptioningModel, self).__init__()
        self.image_encoder = ImageEncoder(embed_size)
        self.text_decoder = TransformerDecoder(embed_size, heads, num_layers, vocab_size)

    def forward(self, images, captions):
        features = self.image_encoder(images)
        output = self.text_decoder(features, captions)
        return output




# 我这里假设构造一个图像特征是一个且对应描述也是一个
vocab_size = 11  # 字典大小，也是后面概率预测数，一般是32000
vocab = {0:"pad",1: "start", 2: "是", 3: '我', 4: "描述", 5: "图像", 6: "模型", 7: "s", 8: "1", 9: "等", 10: "end"}  # 字典大小为9
images_rand = torch.randn(1, 2048)  # 随机生成一张图像特征
caption = torch.tensor([0, 2, 1, 4, 3, 5, 9]).reshape(1,-1)  # 我是图像描述模型
start_vocab = 1
end_vocab = 10
max_length = 20  # 图像长度





# 定义随机生成数据的 Dataset,我将其固定位一个图像
class RandomDataset(Dataset):
    def __init__(self, num_samples, max_length):
        # self.images = torch.randn(num_samples, 2048)  # 随机生成图像特征
        self.images = images_rand  # 随机生成图像特征

        self.captions = F.pad(caption, (0, max_length-max(caption.shape)), value=0)

        self.num_samples = num_samples
    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # 因为我只想使用一张图与一个描述，因此idx只起到占位
        return self.images[0], self.captions[0]




# 推理
def inference(image,model):

    features = model.image_encoder(image)

    start_token = torch.tensor([[start_vocab]])  # 起始标记
    result_caption = [0]  # 存储生成的描述

    for _ in range(20):  # 最大生成长度为20
        output = model.text_decoder(features, start_token)
        output = output.argmax(2)
        word_index = output[0][-1].item()
        result_caption.append(word_index)

        if word_index == end_vocab:  # 到达终止标记则停止生成
            break

        start_token = torch.tensor([[word_index]])

    return result_caption



if __name__ == '__main__':

    # 创建模型和数据加载器
    model = ImageCaptioningModel(embed_size=256, heads=8, num_layers=2, vocab_size=vocab_size)
    dataset = RandomDataset(num_samples=40, max_length=max_length)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    for epoch in range(60):
        for images, captions in dataloader:
            optimizer.zero_grad()
            features = model.image_encoder(images)
            outputs = model.text_decoder(features, captions)

            loss = criterion(outputs.view(-1, vocab_size), captions.view(-1))  # 计算损失
            loss.backward()
            optimizer.step()

            print(f'Epoch {epoch + 1}, Loss: {loss.item()}')



    # 推理
    predicted_caption = inference(images_rand,model)  # images的shape是[1,256]

    print('Predicted Caption:', predicted_caption)
    result = ''.join(vocab[int(index)] for index in predicted_caption)
    print('Predicted Caption:', result)