utils/models.py

from typing import Tuple, Union
from torch import nn
import torch
from .model_utils import VisualTransformer
from .configuration_bert import BertConfig
from .modeling_bert import BertModel
import numpy as np
from utils import _tokenizer

class CLIP(nn.Module):
    def __init__(self,
                 embed_dim: int,
                 # vision
                 image_resolution: int,
                 vision_layers: Union[Tuple[int, int, int, int], int],
                 vision_width: int,
                 vision_patch_size: int,
                 # text
                 vocab_size: int,
                 text_attention_probs_dropout_prob: float,
                 text_hidden_act: str,
                 text_hidden_dropout_prob: float,
                 text_hidden_size: int,
                 text_initializer_range: float,
                 text_intermediate_size: int,
                 text_max_position_embeddings: int,
                 text_num_attention_heads: int,
                 text_num_hidden_layers: int,
                 text_type_vocab_size: int,
                 tokenizer=_tokenizer,
                 ):
        super().__init__()

        vision_heads = vision_width // 64
        self.visual = VisualTransformer(
            input_resolution=image_resolution,
            patch_size=vision_patch_size,
            width=vision_width,
            layers=vision_layers,
            heads=vision_heads,
            output_dim=embed_dim
        )

        self.bert_config = BertConfig(
            vocab_size_or_config_json_file=vocab_size,
            hidden_size=text_hidden_size,
            num_hidden_layers=text_num_hidden_layers,
            num_attention_heads=text_num_attention_heads,
            intermediate_size=text_intermediate_size,
            hidden_act=text_hidden_act,
            hidden_dropout_prob=text_hidden_dropout_prob,
            attention_probs_dropout_prob=text_attention_probs_dropout_prob,
            max_position_embeddings=text_max_position_embeddings,
            type_vocab_size=text_type_vocab_size,
            initializer_range=text_initializer_range,
            layer_norm_eps=1e-12,
        )
        self.bert = BertModel(self.bert_config)

        self.text_projection = nn.Parameter(
            torch.empty(text_hidden_size, embed_dim))
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

        self.tokenizer = tokenizer

        # loss
        # self.cl_head = CLIPLoss_withMask_withmultimodal()

        self.initialize_parameters()

    def initialize_parameters(self):
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

        if self.text_projection is not None:
            nn.init.normal_(self.text_projection,
                            std=self.bert_config.hidden_size ** -0.5)

    @property
    def dtype(self):
        return self.visual.conv1.weight.dtype

    def encode_image(self, image):
        return self.visual(image.type(self.dtype))

    def encode_text(self, text):
        pad_index = self.tokenizer.vocab['[PAD]']
        attn_mask = text.ne(pad_index).type(self.dtype)
        x = self.bert(text, attention_mask=attn_mask)[0].type(
            self.dtype)  # [batch_size, seq_length, hidden_size]
        x = x @ self.text_projection
        return x[:, 0, :], x

    def forward(self, image, text):
        assert image is not None or text is not None, "text and image cannot both be None!"

        if image is None:
            return self.encode_text(text)
        elif text is None:
            return self.encode_image(image)
        image_features = self.encode_image(image)
        text_features = self.encode_text(text)

        features = {'image_embed': image_features,
                    'text_embed': text_features,
                    'logit_scale': self.logit_scale.exp()}

        ret = self.cl_head(features)

        return ret
    


class MergeLayer(nn.Module):
    def __init__(self, d_model, nhead):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=nhead, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model))
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, query, visual_feats, text_feats, key_padding_mask=None):
        visual_cls = visual_feats[:, 0]
        text_cls = text_feats[:, 0]
        query = query + visual_cls.unsqueeze(1) + text_cls.unsqueeze(1)
        kv = torch.cat([visual_feats, text_feats], dim=1)

        # 交叉注意力
        attn_output, _ = self.cross_attn(query=query, key=kv, value=kv, key_padding_mask=key_padding_mask)
        # 残差连接+层归一化
        query = self.norm1(query + attn_output)
        # 前馈网络
        ffn_output = self.ffn(query)
        query = self.norm2(query + ffn_output)
        return query


class Modality_Mergerv3(nn.Module):
    def __init__(self, d_model: int = 512, d_output: int = 256, nhead: int = 8, layer_num: int = 3):
        super().__init__()
        self.q = nn.Parameter(torch.randn(1, 1, d_model))
        self.layers = nn.ModuleList([
            MergeLayer(d_model, nhead) for _ in range(layer_num)
        ])
        self.proj_d512 = nn.Linear(d_model, 512)
        self.proj_d256 = nn.Linear(d_model, 256)
        self.proj_d128 = nn.Linear(d_model, 128)
        self.proj_d64 = nn.Linear(d_model, 64)
        self.proj_d32 = nn.Linear(d_model, 32)

    def forward(self, visual_feats,
                text_feats, text_mask,
                cate_feats, cate_mask,
                c2c_feats,
                ):
        bs = visual_feats.size(0)

        # 创建视觉特征的全1 mask（假设视觉特征无pad）
        visual_mask = torch.ones(visual_feats.size()[:2], dtype=text_mask.dtype, device=text_mask.device)
        c2c_mask = torch.ones(c2c_feats.size()[:2], dtype=text_mask.dtype, device=text_mask.device)
        key_padding_mask = torch.cat([visual_mask, text_mask, cate_mask, c2c_mask], dim=1)
        key_padding_mask = key_padding_mask == 0

        text_feats = torch.cat([text_feats, cate_feats, c2c_feats], dim=1)

        query = self.q.expand(bs, -1, -1)
        for layer in self.layers:
            query = layer(query, visual_feats, text_feats, key_padding_mask)

        query = query.squeeze(1)

        mm_features_d512 = self.proj_d512(query)
        mm_features_d256 = self.proj_d256(query)
        mm_features_d128 = self.proj_d128(query)
        mm_features_d64 = self.proj_d64(query)
        mm_features_d32 = self.proj_d32(query)
        return mm_features_d512, mm_features_d256, mm_features_d128, mm_features_d64, mm_features_d32