331 lines
11 KiB
Python
331 lines
11 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2021 The OpenAI Team Authors.
|
|
# Copyright 2021-2022 The Alibaba Fundamental Vision Team Authors. All rights reserved.
|
|
#
|
|
# The implementation here is modified based on OpenAI CLIP,
|
|
# originally MIT License, Copyright (c) 2021 OpenAI,
|
|
# and publicly available at https://github.com/openai/CLIP/.
|
|
|
|
from collections import OrderedDict
|
|
from typing import Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.utils.checkpoint as checkpoint
|
|
from torch import nn
|
|
from transformers import BertConfig, BertForMaskedLM
|
|
|
|
|
|
class LayerNorm(nn.LayerNorm):
|
|
"""Subclass torch's LayerNorm to handle fp16."""
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
orig_type = x.dtype
|
|
ret = super().forward(x.type(torch.float32))
|
|
return ret.type(orig_type)
|
|
|
|
|
|
class QuickGELU(nn.Module):
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return x * torch.sigmoid(1.702 * x)
|
|
|
|
|
|
class ResidualAttentionBlock(nn.Module):
|
|
|
|
def __init__(self,
|
|
d_model: int,
|
|
n_head: int,
|
|
attn_mask: torch.Tensor = None):
|
|
super().__init__()
|
|
|
|
self.attn = nn.MultiheadAttention(d_model, n_head)
|
|
self.ln_1 = LayerNorm(d_model)
|
|
self.mlp = nn.Sequential(
|
|
OrderedDict([('c_fc', nn.Linear(d_model, d_model * 4)),
|
|
('gelu', QuickGELU()),
|
|
('c_proj', nn.Linear(d_model * 4, d_model))]))
|
|
self.ln_2 = LayerNorm(d_model)
|
|
self.attn_mask = attn_mask
|
|
|
|
def attention(self, x: torch.Tensor):
|
|
self.attn_mask = self.attn_mask.to(
|
|
dtype=x.dtype,
|
|
device=x.device) if self.attn_mask is not None else None
|
|
return self.attn(
|
|
x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
x = x + self.attention(self.ln_1(x))
|
|
x = x + self.mlp(self.ln_2(x))
|
|
return x
|
|
|
|
|
|
class Transformer(nn.Module):
|
|
|
|
def __init__(self,
|
|
width: int,
|
|
layers: int,
|
|
heads: int,
|
|
attn_mask: torch.Tensor = None,
|
|
use_gc=False):
|
|
super().__init__()
|
|
self.use_gc = use_gc
|
|
self.width = width
|
|
self.layers = layers
|
|
self.resblocks = nn.Sequential(*[
|
|
ResidualAttentionBlock(width, heads, attn_mask)
|
|
for _ in range(layers)
|
|
])
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
if self.use_gc:
|
|
for each_block in self.resblocks:
|
|
x = checkpoint.checkpoint(each_block, x)
|
|
return x
|
|
else:
|
|
return self.resblocks(x)
|
|
|
|
|
|
class VisionTransformer(nn.Module):
|
|
|
|
def __init__(self,
|
|
input_resolution: int,
|
|
patch_size: int,
|
|
width: int,
|
|
layers: int,
|
|
heads: int,
|
|
output_dim: int,
|
|
use_gc=False):
|
|
super().__init__()
|
|
self.input_resolution = input_resolution
|
|
self.output_dim = output_dim
|
|
self.conv1 = nn.Conv2d(
|
|
in_channels=3,
|
|
out_channels=width,
|
|
kernel_size=patch_size,
|
|
stride=patch_size,
|
|
bias=False)
|
|
|
|
scale = width**-0.5
|
|
self.class_embedding = nn.Parameter(scale * torch.randn(width))
|
|
self.positional_embedding = nn.Parameter(scale * torch.randn(
|
|
(input_resolution // patch_size)**2 + 1, width))
|
|
self.ln_pre = LayerNorm(width)
|
|
|
|
self.transformer = Transformer(width, layers, heads, use_gc=use_gc)
|
|
|
|
self.ln_post = LayerNorm(width)
|
|
self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
x = self.conv1(x) # shape = [*, width, grid, grid]
|
|
x = x.reshape(x.shape[0], x.shape[1],
|
|
-1) # shape = [*, width, grid ** 2]
|
|
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
|
|
class_embedding = self.class_embedding.to(x.dtype) + \
|
|
torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device)
|
|
x = torch.cat([class_embedding, x],
|
|
dim=1) # shape = [*, grid ** 2 + 1, width]
|
|
x = x + self.positional_embedding.to(x.dtype)
|
|
x = self.ln_pre(x)
|
|
|
|
x = x.permute(1, 0, 2) # NLD -> LND
|
|
x = self.transformer(x)
|
|
x = x.permute(1, 0, 2) # LND -> NLD
|
|
|
|
x = self.ln_post(x[:, 0, :])
|
|
|
|
if self.proj is not None:
|
|
x = x @ self.proj
|
|
|
|
return x
|
|
|
|
|
|
class CLIPVisionWrapper(nn.Module):
|
|
|
|
def __init__(self, ):
|
|
super().__init__()
|
|
self.vision_transformer = VisionTransformer(
|
|
input_resolution=224,
|
|
patch_size=14,
|
|
width=1024,
|
|
layers=24,
|
|
heads=16,
|
|
output_dim=768)
|
|
|
|
def forward(self, x):
|
|
x = self.vision_transformer.conv1(x) # shape = [*, width, grid, grid]
|
|
x = x.reshape(x.shape[0], x.shape[1],
|
|
-1) # shape = [*, width, grid ** 2]
|
|
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
|
|
class_embedding = self.vision_transformer.class_embedding.to(x.dtype) + \
|
|
torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device)
|
|
x = torch.cat([class_embedding, x],
|
|
dim=1) # shape = [*, grid ** 2 + 1, width]
|
|
x = x + self.vision_transformer.positional_embedding.to(x.dtype)
|
|
x = self.vision_transformer.ln_pre(x)
|
|
|
|
x = x.permute(1, 0, 2) # NLD -> LND
|
|
x = self.vision_transformer.transformer(x)
|
|
x = x.permute(1, 0, 2) # LND -> NLD
|
|
|
|
x_tensor = x.clone()
|
|
x = self.vision_transformer.ln_post(x[:, 0, :])
|
|
|
|
if self.vision_transformer.proj is not None:
|
|
x = x @ self.vision_transformer.proj
|
|
|
|
return x, x_tensor
|
|
|
|
|
|
class BertWrapper(nn.Module):
|
|
|
|
def __init__(self, config_json, feat_dim, token_dim):
|
|
super(BertWrapper, self).__init__()
|
|
bert_config = BertConfig.from_json_file(config_json)
|
|
self.bert = BertForMaskedLM(bert_config).bert
|
|
|
|
self.projector = nn.Linear(768, feat_dim, bias=False)
|
|
self.projector_token_embeds = nn.Linear(768, token_dim)
|
|
|
|
def forward(self, input_ids, attention_mask):
|
|
trans_features = {
|
|
'input_ids': input_ids,
|
|
'attention_mask': attention_mask
|
|
}
|
|
output_states = self.bert(**trans_features, return_dict=False)
|
|
output_tokens = output_states[0]
|
|
|
|
cls_tokens = output_tokens[:, 0, :] # CLS token is first token
|
|
|
|
return self.projector(cls_tokens), self.projector_token_embeds(
|
|
output_tokens)
|
|
|
|
|
|
class Mlp(nn.Module):
|
|
|
|
def __init__(self,
|
|
in_features,
|
|
hidden_features=None,
|
|
out_features=None,
|
|
act_layer=nn.GELU,
|
|
drop=0.):
|
|
super().__init__()
|
|
out_features = out_features or in_features
|
|
hidden_features = hidden_features or in_features
|
|
self.fc1 = nn.Linear(in_features, hidden_features)
|
|
self.act = act_layer()
|
|
self.fc2 = nn.Linear(hidden_features, out_features)
|
|
self.drop = nn.Dropout(drop)
|
|
|
|
def forward(self, x):
|
|
x = self.fc1(x)
|
|
x = self.act(x)
|
|
x = self.drop(x)
|
|
x = self.fc2(x)
|
|
x = self.drop(x)
|
|
return x
|
|
|
|
|
|
class CrossLayer(nn.Module):
|
|
|
|
def __init__(self, feat_dim, mlp_ratio):
|
|
super(CrossLayer, self).__init__()
|
|
self.norm1 = nn.LayerNorm(feat_dim)
|
|
self.norm2 = nn.LayerNorm(feat_dim)
|
|
self.norm3 = nn.LayerNorm(feat_dim)
|
|
|
|
self.self_attn = nn.MultiheadAttention(
|
|
embed_dim=feat_dim, num_heads=16)
|
|
self.cross_attn = nn.MultiheadAttention(
|
|
embed_dim=feat_dim, num_heads=16)
|
|
self.ffn = Mlp(
|
|
in_features=feat_dim,
|
|
hidden_features=feat_dim * mlp_ratio,
|
|
drop=0.1)
|
|
|
|
self.dropout1 = nn.Dropout(0.1)
|
|
self.dropout2 = nn.Dropout(0.1)
|
|
self.dropout3 = nn.Dropout(0.1)
|
|
|
|
def forward(self, text_tensors, text_masks, image_tensors,
|
|
retrieved_tensors):
|
|
retrieved_tensors_res = self.norm1(retrieved_tensors)
|
|
retrieved_tensors_res = self.self_attn(
|
|
(text_tensors + retrieved_tensors_res).permute(1, 0, 2),
|
|
(text_tensors + retrieved_tensors_res).permute(1, 0, 2),
|
|
retrieved_tensors_res.permute(1, 0, 2),
|
|
key_padding_mask=(text_masks == 0),
|
|
)[0].permute(1, 0, 2)
|
|
retrieved_tensors = retrieved_tensors + self.dropout1(
|
|
retrieved_tensors_res)
|
|
|
|
retrieved_tensors_res = self.norm2(retrieved_tensors)
|
|
retrieved_tensors_res = self.cross_attn(
|
|
(text_tensors + retrieved_tensors_res).permute(1, 0, 2),
|
|
image_tensors.permute(1, 0, 2),
|
|
image_tensors.permute(1, 0, 2))[0].permute(1, 0, 2)
|
|
retrieved_tensors = retrieved_tensors + self.dropout2(
|
|
retrieved_tensors_res)
|
|
|
|
retrieved_tensors_res = self.norm3(retrieved_tensors)
|
|
retrieved_tensors = retrieved_tensors + self.dropout3(
|
|
self.ffn(retrieved_tensors_res))
|
|
|
|
return retrieved_tensors
|
|
|
|
|
|
class TEAM(nn.Module):
|
|
|
|
def __init__(self, text_model, image_model, pretrained):
|
|
super(TEAM, self).__init__()
|
|
self.text_model = text_model
|
|
self.image_model = image_model
|
|
|
|
self.cross_model = nn.ModuleList(
|
|
[CrossLayer(feat_dim=1024, mlp_ratio=2)])
|
|
|
|
self.image_tensor_fc = nn.Linear(1024, 768)
|
|
self.text_tensor_fc = nn.Linear(1024, 768)
|
|
|
|
params = torch.load(pretrained, 'cpu')
|
|
self.load_state_dict(params, strict=True)
|
|
|
|
def get_feature(self, text_data=None, text_mask=None, img_tensor=None):
|
|
if text_data is not None:
|
|
text_feature, text_tensors = self.text_model(text_data, text_mask)
|
|
text_feature = F.normalize(text_feature, p=2.0, dim=1)
|
|
else:
|
|
text_feature, text_tensors = None, None
|
|
|
|
if img_tensor is not None:
|
|
image_feature, image_tensors = self.image_model(img_tensor)
|
|
image_feature = F.normalize(image_feature, p=2.0, dim=1)
|
|
else:
|
|
image_feature, image_tensors = None, None
|
|
|
|
return text_feature, text_tensors, image_feature, image_tensors
|
|
|
|
def get_cross_score(self, text_tensors, text_mask, image_tensors):
|
|
retrieved_tensors = torch.zeros_like(text_tensors)
|
|
pair_score_list = []
|
|
text_tensors_proj = self.text_tensor_fc(text_tensors)
|
|
text_mask_float = text_mask.type(text_tensors_proj.dtype)
|
|
for each_cross_model in self.cross_model:
|
|
retrieved_tensors = each_cross_model(text_tensors, text_mask,
|
|
image_tensors,
|
|
retrieved_tensors)
|
|
retrieved_tensors_proj = self.image_tensor_fc(retrieved_tensors)
|
|
|
|
pair_score = torch.sum(
|
|
F.normalize(retrieved_tensors_proj, p=2.0, dim=2)
|
|
* F.normalize(text_tensors_proj, p=2.0, dim=2),
|
|
dim=2)
|
|
pair_score_reduced = torch.sum(
|
|
pair_score * text_mask_float, dim=1) / torch.clamp(
|
|
torch.sum(text_mask_float, dim=1), min=1.0)
|
|
pair_score_list.append(pair_score_reduced)
|
|
return pair_score_list
|