560 lines
21 KiB
Python
560 lines
21 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright 2021 The OpenAI Team Authors.
|
|
# Copyright 2022 Phil Wang.
|
|
# 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/.
|
|
# The implementation here is modified based on Coca-pytorch,
|
|
# originally MIT License, Copyright (c) 2022 Phil Wang,
|
|
# and publicly available at https://github.com/lucidrains/CoCa-pytorch/,
|
|
""" Generative Multimodal Model Architecture."""
|
|
|
|
import os
|
|
from collections import OrderedDict
|
|
from typing import Tuple, Union
|
|
|
|
import json
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
from torch.nn import LayerNorm
|
|
|
|
from modelscope.models.multi_modal.gemm.tokenizer import (SimpleTokenizer,
|
|
clip_tokenize)
|
|
|
|
|
|
class Bottleneck(nn.Module):
|
|
""" ResNet style bottleneck module
|
|
From https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
expansion = 4
|
|
|
|
def __init__(self, inplanes, planes, stride=1):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(planes)
|
|
self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(planes)
|
|
self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
|
|
self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
|
|
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self.downsample = None
|
|
self.stride = stride
|
|
if stride > 1 or inplanes != planes * Bottleneck.expansion:
|
|
self.downsample = nn.Sequential(
|
|
OrderedDict([('-1', nn.AvgPool2d(stride)),
|
|
('0',
|
|
nn.Conv2d(
|
|
inplanes,
|
|
planes * self.expansion,
|
|
1,
|
|
stride=1,
|
|
bias=False)),
|
|
('1', nn.BatchNorm2d(planes * self.expansion))]))
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
identity = x
|
|
out = self.relu(self.bn1(self.conv1(x)))
|
|
out = self.relu(self.bn2(self.conv2(out)))
|
|
out = self.avgpool(out)
|
|
out = self.bn3(self.conv3(out))
|
|
if self.downsample is not None:
|
|
identity = self.downsample(x)
|
|
out += identity
|
|
out = self.relu(out)
|
|
return out
|
|
|
|
|
|
class QuickGELU(nn.Module):
|
|
""" A quick version of GELU module
|
|
From https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return x * torch.sigmoid(1.702 * x)
|
|
|
|
|
|
class ResidualAttentionBlock(nn.Module):
|
|
""" Multihead attention block with residual link
|
|
Adapted from https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
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
|
|
attn_mask = self.attn_mask
|
|
if attn_mask is not None and attn_mask.shape[0] > x.shape[0]:
|
|
attn_mask = self.attn_mask[:x.shape[0], :x.shape[0]]
|
|
return self.attn(x, x, x, need_weights=False, attn_mask=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):
|
|
""" Transformer encoder module
|
|
Adapted from https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
def __init__(self,
|
|
width: int,
|
|
layers: int,
|
|
heads: int,
|
|
attn_mask: torch.Tensor = None,
|
|
use_gc: bool = 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):
|
|
return self.resblocks(x)
|
|
|
|
|
|
class AttentionPool2d(nn.Module):
|
|
""" Pool layer with attention module
|
|
Adapted from https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
def __init__(self,
|
|
spacial_dim: int,
|
|
embed_dim: int,
|
|
num_heads: int,
|
|
output_dim: int = None):
|
|
super().__init__()
|
|
self.positional_embedding = nn.Parameter(
|
|
torch.randn(spacial_dim**2 + 1, embed_dim) / embed_dim**0.5)
|
|
self.k_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.q_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.v_proj = nn.Linear(embed_dim, embed_dim)
|
|
self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
|
|
self.num_heads = num_heads
|
|
|
|
def forward(self, x):
|
|
x = x.reshape(x.shape[0], x.shape[1],
|
|
x.shape[2] * x.shape[3]).permute(2, 0, 1)
|
|
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)
|
|
x = x + self.positional_embedding[:, None, :].to(x.dtype)
|
|
x, _ = F.multi_head_attention_forward(
|
|
query=x,
|
|
key=x,
|
|
value=x,
|
|
embed_dim_to_check=x.shape[-1],
|
|
num_heads=self.num_heads,
|
|
q_proj_weight=self.q_proj.weight,
|
|
k_proj_weight=self.k_proj.weight,
|
|
v_proj_weight=self.v_proj.weight,
|
|
in_proj_weight=None,
|
|
in_proj_bias=torch.cat(
|
|
[self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
|
|
bias_k=None,
|
|
bias_v=None,
|
|
add_zero_attn=False,
|
|
dropout_p=0,
|
|
out_proj_weight=self.c_proj.weight,
|
|
out_proj_bias=self.c_proj.bias,
|
|
use_separate_proj_weight=True,
|
|
training=self.training,
|
|
need_weights=False)
|
|
return x.permute(1, 0, 2).contiguous()
|
|
|
|
|
|
class CrossAttention(nn.Module):
|
|
""" Cross attention module with query and context as input
|
|
Adapted from https://github.com/lucidrains/CoCa-pytorch/blob/main/coca_pytorch/coca_pytorch.py
|
|
"""
|
|
|
|
def __init__(self,
|
|
dim,
|
|
*,
|
|
context_dim=None,
|
|
dim_head=64,
|
|
heads=8,
|
|
parallel_ff=False,
|
|
ff_mult=4,
|
|
norm_context=False):
|
|
super().__init__()
|
|
self.heads = heads
|
|
self.scale = dim_head**-0.5
|
|
inner_dim = heads * dim_head
|
|
context_dim = dim if context_dim is None else context_dim
|
|
self.norm = LayerNorm(dim)
|
|
self.context_norm = LayerNorm(
|
|
context_dim) if norm_context else nn.Identity()
|
|
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
|
self.to_kv = nn.Linear(context_dim, dim_head * 2, bias=False)
|
|
self.to_out = nn.Linear(inner_dim, dim, bias=False)
|
|
ff_inner_dim = ff_mult * dim
|
|
self.ff = nn.Sequential(
|
|
nn.Linear(dim, ff_inner_dim * 2, bias=False), SwiGLU(),
|
|
nn.Linear(ff_inner_dim, dim, bias=False)) if parallel_ff else None
|
|
|
|
def forward(self, x, context):
|
|
"""
|
|
einstein notation
|
|
b - batch
|
|
h - heads
|
|
n, i, j - sequence length (base sequence length, source, target)
|
|
d - feature dimension
|
|
"""
|
|
|
|
x = self.norm(x)
|
|
context = self.context_norm(context)
|
|
|
|
q = self.to_q(x)
|
|
q = q.view(q.shape[0], q.shape[1], self.heads,
|
|
-1).permute(0, 2, 1, 3).contiguous()
|
|
q = q * self.scale
|
|
k, v = self.to_kv(context).chunk(2, dim=-1)
|
|
sim = torch.einsum('b h i d, b j d -> b h i j', q, k)
|
|
sim = sim - sim.amax(dim=-1, keepdim=True)
|
|
attn = sim.softmax(dim=-1)
|
|
out = torch.einsum('b h i j, b j d -> b h i d', attn, v)
|
|
out = out.permute(0, 2, 1,
|
|
3).contiguous().reshape(out.shape[0], out.shape[2],
|
|
-1)
|
|
out = self.to_out(out)
|
|
if self.ff is not None:
|
|
out = out + self.ff(x)
|
|
return out
|
|
|
|
|
|
class ModifiedResNet(nn.Module):
|
|
""" Modified ResNet backbone
|
|
From https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
A ResNet class that is similar to torchvision's but contains the following changes:
|
|
- There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
|
|
- Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
|
|
- The final pooling layer is a QKV attention instead of an average pool
|
|
"""
|
|
|
|
def __init__(self,
|
|
layers,
|
|
output_dim,
|
|
heads,
|
|
input_resolution=224,
|
|
width=64):
|
|
super().__init__()
|
|
self.output_dim = output_dim
|
|
self.input_resolution = input_resolution
|
|
self.conv1 = nn.Conv2d(
|
|
3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
|
|
self.bn1 = nn.BatchNorm2d(width // 2)
|
|
self.conv2 = nn.Conv2d(
|
|
width // 2, width // 2, kernel_size=3, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm2d(width // 2)
|
|
self.conv3 = nn.Conv2d(
|
|
width // 2, width, kernel_size=3, padding=1, bias=False)
|
|
self.bn3 = nn.BatchNorm2d(width)
|
|
self.avgpool = nn.AvgPool2d(2)
|
|
self.relu = nn.ReLU(inplace=True)
|
|
self._inplanes = width
|
|
self.layer1 = self._make_layer(width, layers[0])
|
|
self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
|
|
self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
|
|
self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
|
|
|
|
embed_dim = width * 32
|
|
self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim,
|
|
heads, output_dim)
|
|
|
|
def _make_layer(self, planes, blocks, stride=1):
|
|
layers = [Bottleneck(self._inplanes, planes, stride)]
|
|
|
|
self._inplanes = planes * Bottleneck.expansion
|
|
for _ in range(1, blocks):
|
|
layers.append(Bottleneck(self._inplanes, planes))
|
|
|
|
return nn.Sequential(*layers)
|
|
|
|
def forward(self, x):
|
|
|
|
def stem(x):
|
|
for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2),
|
|
(self.conv3, self.bn3)]:
|
|
x = self.relu(bn(conv(x)))
|
|
x = self.avgpool(x)
|
|
return x
|
|
|
|
x = stem(x)
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
x = self.attnpool(x)
|
|
|
|
return x
|
|
|
|
|
|
class VisualTransformer(nn.Module):
|
|
""" ViT transformer backbone
|
|
From https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
def __init__(self, input_resolution: int, patch_size: int, width: int,
|
|
layers: int, heads: int, output_dim: int, use_gc: bool):
|
|
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)
|
|
x = x.reshape(x.shape[0], x.shape[1], -1)
|
|
x = x.permute(0, 2, 1)
|
|
z = torch.zeros(
|
|
x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device)
|
|
x = torch.cat([self.class_embedding.to(x.dtype) + z, x], dim=1)
|
|
x = x + self.positional_embedding.to(x.dtype)
|
|
x = self.ln_pre(x)
|
|
x = x.permute(1, 0, 2)
|
|
x = self.transformer(x)
|
|
x = x.permute(1, 0, 2)
|
|
x = self.ln_post(x)
|
|
if self.proj is not None:
|
|
x = x @ self.proj
|
|
return x
|
|
|
|
|
|
class GEVL(nn.Module):
|
|
""" Generative vision-language model
|
|
Support learning from both generative and contrastive loss.
|
|
Given image and text input, it could output the features of
|
|
image and text respectively. Furthermore, caption could also
|
|
be produced when image input is available.
|
|
"""
|
|
|
|
def __init__(self, embed_dim: int, image_resolution: int,
|
|
vision_layers: Union[Tuple[int, int, int, int],
|
|
int], vision_width: int,
|
|
vision_patch_size: int, context_length: int, vocab_size: int,
|
|
transformer_width: int, transformer_heads: int,
|
|
transformer_layers: int, use_gc: bool, tokenizer):
|
|
nn.Module.__init__(self)
|
|
self.context_length = context_length
|
|
self.vis_token_size = context_length
|
|
self.tokenizer = tokenizer
|
|
|
|
if isinstance(vision_layers, (tuple, list)):
|
|
vision_heads = vision_width * 32 // 64
|
|
self.visual = ModifiedResNet(
|
|
layers=vision_layers,
|
|
output_dim=embed_dim,
|
|
heads=vision_heads,
|
|
input_resolution=image_resolution,
|
|
width=vision_width)
|
|
else:
|
|
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,
|
|
use_gc=use_gc)
|
|
|
|
self.transformer = Transformer(
|
|
width=transformer_width,
|
|
layers=transformer_layers,
|
|
heads=transformer_heads,
|
|
attn_mask=self.build_attention_mask(),
|
|
use_gc=use_gc)
|
|
|
|
self.vocab_size = vocab_size
|
|
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
|
|
self.positional_embedding = nn.Parameter(
|
|
torch.empty(self.context_length, transformer_width))
|
|
self.ln_final = LayerNorm(transformer_width)
|
|
|
|
self.vis_token_projection = nn.Parameter(
|
|
torch.empty(embed_dim, transformer_width))
|
|
nn.init.normal_(
|
|
self.vis_token_projection, std=self.transformer.width**-0.5)
|
|
self.text_projection = nn.Parameter(
|
|
torch.empty(transformer_width, embed_dim))
|
|
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
|
|
self.decoder = Transformer(
|
|
width=transformer_width,
|
|
layers=4,
|
|
heads=transformer_heads,
|
|
attn_mask=self.build_attention_mask(
|
|
self.vis_token_size + self.context_length,
|
|
self.vis_token_size),
|
|
use_gc=use_gc)
|
|
self.to_logits = nn.Sequential(
|
|
LayerNorm(transformer_width),
|
|
nn.Linear(transformer_width, transformer_width),
|
|
nn.Linear(transformer_width, vocab_size, bias=False))
|
|
self.gen_logit_scale = nn.Parameter(
|
|
torch.ones([]) * np.log(np.log(vocab_size)))
|
|
self.bias = nn.Parameter(torch.ones(vocab_size))
|
|
self.to_logits[-1].weight = self.token_embedding.weight
|
|
self.to_logits[-1].bias = self.bias
|
|
self.img_queries = nn.Parameter(
|
|
torch.randn(self.vis_token_size, transformer_width))
|
|
self.img_attn_pool = CrossAttention(
|
|
dim=transformer_width, norm_context=True)
|
|
self.img_attn_pool_norm = LayerNorm(transformer_width)
|
|
|
|
def build_attention_mask(self, seq_length=None, prefix_length=0):
|
|
seq_length = self.context_length if seq_length is None else seq_length
|
|
mask = torch.empty(seq_length, seq_length)
|
|
mask.fill_(torch.tensor(torch.finfo(torch.float16).min))
|
|
mask.triu_(1)
|
|
if prefix_length > 0:
|
|
mask[:prefix_length, :prefix_length] = 0
|
|
return mask
|
|
|
|
@property
|
|
def dtype(self):
|
|
return self.visual.conv1.weight.dtype
|
|
|
|
def encode_image(self, image, return_tokens=False):
|
|
image_outputs = self.visual(image)
|
|
image_features = image_outputs[:, 0, :]
|
|
image_features = image_features / image_features.norm(
|
|
dim=-1, p=2, keepdim=True)
|
|
if return_tokens:
|
|
image_tokens = image_outputs[:, 1:, :] @ self.vis_token_projection
|
|
return image_features, image_tokens
|
|
else:
|
|
return image_features
|
|
|
|
def encode_text(self, text, return_tokens=False):
|
|
x = self.token_embedding(text)
|
|
x = x + self.positional_embedding[:x.shape[1], :]
|
|
x = x.permute(1, 0, 2)
|
|
x = self.transformer(x)
|
|
x = x.permute(1, 0, 2)
|
|
x = self.ln_final(x)
|
|
text_features = x[torch.arange(x.shape[0]),
|
|
text.argmax(dim=-1), ...] @ self.text_projection
|
|
text_features = text_features / text_features.norm(
|
|
dim=-1, p=2, keepdim=True)
|
|
if return_tokens:
|
|
text_tokens = x
|
|
return text_features, text_tokens
|
|
else:
|
|
return text_features
|
|
|
|
def image_to_text(self, image):
|
|
image_features, image_tokens = self.encode_image(
|
|
image, return_tokens=True)
|
|
img_queries = self.img_queries.expand(image_tokens.shape[0], -1, -1)
|
|
img_token_features = self.img_attn_pool(img_queries, image_tokens)
|
|
img_token_features = self.img_attn_pool_norm(img_token_features)
|
|
sot_token = self.tokenizer.encoder['<|startoftext|>']
|
|
eot_token = self.tokenizer.encoder['<|endoftext|>']
|
|
text_input = image.new_ones(
|
|
image.shape[0], 1, dtype=torch.long) * sot_token
|
|
input_tokens = img_token_features
|
|
pred_tokens = []
|
|
for text_idx in range(self.context_length):
|
|
text_features, text_tokens = self.encode_text(
|
|
text_input, return_tokens=True)
|
|
input_tokens = torch.cat([img_token_features, text_tokens], axis=1)
|
|
out_embs = self.decoder(input_tokens.permute(1, 0, 2).contiguous())
|
|
gen_logits = self.to_logits(out_embs[-1:, ...])
|
|
probs = F.softmax(self.gen_logit_scale.exp() * gen_logits, dim=-1)
|
|
pred = torch.argmax(
|
|
probs * (2.0 + torch.rand_like(probs)), axis=-1)
|
|
if int(pred) >= eot_token or int(pred) <= 0:
|
|
break
|
|
pred_tokens.append(pred)
|
|
text_input = torch.cat(
|
|
[text_input, pred.permute(1, 0).contiguous()], axis=1)
|
|
pred_text_tokens = torch.cat(pred_tokens, axis=0).permute(1, 0)
|
|
text_list = []
|
|
for out_tokens in pred_text_tokens:
|
|
tokens = []
|
|
for x in out_tokens:
|
|
tokens.append(int(x))
|
|
out_text = self.tokenizer.decode(tokens)
|
|
out_text = out_text.strip()
|
|
text_list.append(out_text)
|
|
return image_features, text_list[0]
|
|
|
|
|
|
class GEMMModel(nn.Module):
|
|
""" Generative multi-modal model, wrapper of GEVL module.
|
|
It takes image or text or both of them as input, and output
|
|
features of input or caption when image input is available.
|
|
"""
|
|
|
|
def __init__(self, model_dir):
|
|
super().__init__()
|
|
with open(
|
|
'{}/encoder_config.json'.format(model_dir), 'r',
|
|
encoding='utf-8') as f:
|
|
model_config = json.loads(f.read())
|
|
model_name = list(model_config.keys())[0]
|
|
config_args = model_config[model_name]
|
|
bpe_path = os.path.join(model_dir, 'bpe_vocab_16e6.txt.gz')
|
|
self.tokenizer = SimpleTokenizer(bpe_path)
|
|
self.model = GEVL(*config_args, self.tokenizer)
|
|
|
|
def tokenize(self, text_str):
|
|
text_tensor = clip_tokenize(self.tokenizer, [text_str])[0]
|
|
return text_tensor
|
|
|
|
def parse_feat(self, feat):
|
|
out = feat.cpu().numpy()
|
|
return out
|
|
|
|
@torch.no_grad()
|
|
def forward(self, image=None, text=None, captioning=True):
|
|
img_feature, text_feature, caption = None, None, None
|
|
if captioning and image is not None:
|
|
img_feature, caption = self.model.image_to_text(image)
|
|
img_feature = self.parse_feat(img_feature)
|
|
elif image is not None:
|
|
img_feature = self.parse_feat(self.model.encode_image(image))
|
|
if text is not None:
|
|
text_feature = self.parse_feat(self.model.encode_text(text))
|
|
out = {
|
|
'image_feature': img_feature,
|
|
'text_feature': text_feature,
|
|
'caption': caption,
|
|
}
|
|
return out
|