多模态原理--CLIP

def tokenizer(text, encode=True,max_seq_length=32):
    if encode:
        out = chr(2) + text + chr(3)  # 添加 SOT token 和 EOT token
        out = out + chr(0)*(max_seq_length - len(out))  # 添加 Padding 字符
        out = torch.tensor([ord(c) for c in out])  # 对文本进行编码
        mask = (out>0).to(torch.int)
        # mask 为什么需要形成方阵?
        mask = mask.expand(max_seq_length, max_seq_length)
    else:
        # 将input_ids解码为text文本
        out = "".join([chr(x) for x in text[1:text.index(0)-1]])
        mask = None
    return out, mask

# 文本 Encoder 模型,并将EOT向量映射到维度为 emb_dim 的向量上
class TextEncoder(nn.Module):
    def __init__(self, vocab_size, width, max_seq_length, n_heads, n_layers, emb_dim):
        super().__init__()
        self.max_seq_length = max_seq_length
        self.encoder_embedding = nn.Embedding(vocab_size, width)
        self.positional_embedding = PositionalEncoding(width, max_seq_length)
        self.encoder = nn.ModuleList([TransformerEncoder(width, n_heads) for _ in range(n_layers)])
        # 可学习投影（projection）
        self.projection = nn.Linear(width, emb_dim, bias=False)

    def forward(self, text, mask):
        # 文本嵌入
        x = self.encoder_embedding(text)
        # 位置嵌入
        x = self.positional_embedding(x)
        # Transformer编码器
        for encoder_layer in self.encoder:
            x = encoder_layer(x, mask=mask)
        # 从EOT的嵌入抽取特征
        x = x[
            torch.arange(text.shape[0]),  # 批次中数据的索引
            torch.sub(torch.sum(mask[:, 0], dim=1), 1)  # 取出掩码mask矩阵的第0行，加和再减1，就得到了EOT的索引
        ]
        # 将文本特征嵌入到联合嵌入空间中（多模态嵌入空间）
        # 文本编码器输出的张量的维度和图像编码器输出的张量的维度必须一致
        if self.projection is not None:
            x = self.projection(x)
        # 除以向量的模长或者范式
        x = x / torch.norm(x, dim=-1, keepdim=True)
        return x

# 图像 Encoder 模型,并将 cls_token 向量映射到维度为 emb_dim 的向量上
class ImageEncoder(nn.Module):
    def __init__(self, width, img_size, patch_size, n_channels, n_layers, n_heads, emb_dim):
        super().__init__()

        assert img_size[0] % patch_size[0] == 0 \
               and img_size[1] % patch_size[1] == 0, \
            "img_size必须能被patch_size整除"
        assert width % n_heads == 0, \
            "width必须能被n_heads整除"

        self.n_patches = (img_size[0] * img_size[1]) // (patch_size[0] * patch_size[1])
        self.max_seq_length = self.n_patches + 1
        self.linear_project = nn.Conv2d(n_channels, width, kernel_size=patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.randn(1, 1, width))
        self.positional_embedding = PositionalEncoding(width, self.max_seq_length)
        self.encoder = nn.ModuleList([
            TransformerEncoder(width, n_heads)
            for _ in range(n_layers)
        ])
        # 可学习的投影
        self.projection = nn.Linear(width, emb_dim, bias=False)

    def forward(self, x):
        # 补丁嵌入
        x = self.linear_project(x)
        x = x.flatten(2).transpose(1, 2)

        # 位置嵌入
        x = torch.cat((self.cls_token.expand(x.size()[0], -1, -1), x), dim=1)
        x = self.positional_embedding(x)

        # Transformer编码器
        for encoder_layer in self.encoder:
            x = encoder_layer(x)

        # 获取类别token
        x = x[:, 0, :]

        # 多模态嵌入
        # 保证文本编码器的输出的维度和图像编码器的输出的维度相等
        if self.projection is not None:
            x = self.projection(x)
        x = x / torch.norm(x, dim=-1, keepdim=True)
        return x

class CLIP(nn.Module):
    def __init__(self, emb_dim, vit_width, img_size, patch_size, n_channels, vit_layers,
                 vit_heads, vocab_size, text_width, max_seq_length, text_heads, text_layers):
        super().__init__()
        self.image_encoder = ImageEncoder(vit_width, img_size, patch_size, n_channels, vit_layers, vit_heads, emb_dim)
        self.text_encoder = TextEncoder(vocab_size, text_width, max_seq_length, text_heads, text_layers, emb_dim)
        # 可学习温度
        self.temperature = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

    def forward(self, image, text, mask):
        # Iₑ是图像嵌入，形状 [B, D=emb_dim]
        I_e = self.image_encoder(image)
        # Tₑ是文本嵌入，形状 [B, D=emb_dim]
        T_e = self.text_encoder(text, mask=mask)

        # 缩放逐点余弦相似度[n, n]
        # 形状 I_e @ T_e^T : [B, D] @ [D, B] --> [B, B]
        logits = (I_e @ T_e.transpose(-2, -1)) * torch.exp(self.temperature)

        # 对称损失函数 labels形状为[B]，值为 [0, 1, 2, ..., B-1]
        labels = torch.arange(logits.shape[0]).to(device)
        # 从文本 --> 图像方向，以文本嵌入 T₃ 为例子，
        # 交叉熵损失的目标是让 T₃⋅I₃ 越大越好
        loss_i = nn.functional.cross_entropy(logits.transpose(-2, -1), labels)
        # 从图像 --> 文本方向，以图像嵌入 I₃ 为例子，
        # 交叉熵损失的目标是让 I₃⋅T₃ 越大越好
        loss_t = nn.functional.cross_entropy(logits, labels)
        # 两个方向的损失求平均值
        loss = (loss_i + loss_t) / 2
        return loss

# 基础配置
ROOT_DIR = Path(__file__).parent.parent
device = 'cuda' if torch.cuda.is_available() else 'cpu'
log_dir = ROOT_DIR / 'logs'

# 超参数配置
emb_dim = 32  # 文本编码器和图像编码器输出的张量的维度
vit_width = 9  # 图像编码器的嵌入的宽度
img_size = (28, 28)
patch_size = (14, 14)
n_channels = 1
vit_layers = 3  # 图像编码器中编码器的层的数量
vit_heads = 3  # 图像编码器中注意力头的数量
vocab_size = 256  # 词汇表大小
text_width = 32  # 文本编码器的嵌入的宽度
max_seq_length = 32  # 最大序列长度
text_heads = 8  # 文本编码器中注意力头的数量
text_layers = 4  # 文本编码器中编码器的层数
lr = 1e-3  # 学习率
epochs = 10
batch_size = 128

# 图片 label 和文本对应关系
captions_dict = {
    0: "An image of 0",
    1: "An image of 1",
    2: "An image of 2",
    3: "An image of 3",
    4: "An image of 4",
    5: "An image of 5",
    6: "An image of 6",
    7: "An image of 7",
    8: "An image of 8",
    9: "An image of 9"
}

# 加载数据
train_set = HandWritingMNIST(train=True, captions_map=captions_dict)
test_set = HandWritingMNIST(train=False, captions_map=captions_dict)

# 数据分批
train_loader = DataLoader(train_set, shuffle=True, batch_size=batch_size)
test_loader = DataLoader(test_set, shuffle=False, batch_size=batch_size)

# 模型初始化
model = CLIP(emb_dim, vit_width, img_size, patch_size, n_channels, vit_layers, vit_heads,
             vocab_size, text_width, max_seq_length, text_heads, text_layers).to(device)

# 优化器
optimizer = optim.Adam(model.parameters(), lr=lr)

best_loss = np.inf
# 开始训练
with SummaryWriter(log_dir=str(log_dir / time.strftime('%Y-%m-%d_%H-%M-%S'))) as writer:
    for epoch in range(epochs):
        for img, _, cap, mask in train_loader:
            img, cap, mask = img.to(device), cap.to(device), mask.to(device)
            loss = model(img, cap, mask)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch [{epoch + 1}/{epochs}], Batch Loss: {loss.item():.3f}")
        # 保存模型
        if loss.item() <= best_loss:
            best_loss = loss.item()
            torch.save(model.state_dict(), "./clip.pt")
            print("模型已经保存...")
        writer.add_scalar('loss', loss.item(), epoch + 1)