aleo1's picture
Upload 41 files
bb6012a verified
raw
history blame
50 kB
from collections import OrderedDict
from typing import Tuple, Union
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from ..utils.dataset import tokenize
from ..utils.simple_tokenizer import SimpleTokenizer as _Tokenizer
_tokenizer = _Tokenizer()
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1):
super().__init__()
# all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
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:
# downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
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
"""
attenpool used in CRIS (output: C1/C2/C3 3 deiffent feature maps)
"""
class ModifiedAttentionPool2d(nn.Module):
def __init__(self,
spacial_dim: int,
embed_dim: int,
num_heads: int,
output_dim: int = None):
super().__init__()
self.spacial_dim = spacial_dim
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
# residual
self.connect = nn.Sequential(
nn.Conv2d(embed_dim, output_dim, 1, stride=1, bias=False),
nn.BatchNorm2d(output_dim))
def resize_pos_embed(self, pos_embed, input_shpae):
"""Resize pos_embed weights.
Resize pos_embed using bicubic interpolate method.
Args:
pos_embed (torch.Tensor): Position embedding weights.
input_shpae (tuple): Tuple for (downsampled input image height,
downsampled input image width).
pos_shape (tuple): The resolution of downsampled origin training
image.
mode (str): Algorithm used for upsampling:
``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
``'trilinear'``. Default: ``'nearest'``
Return:
torch.Tensor: The resized pos_embed of shape [B, C, L_new]
"""
assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]'
pos_h = pos_w = self.spacial_dim
cls_token_weight = pos_embed[:, 0]
pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
pos_embed_weight = pos_embed_weight.reshape(
1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
pos_embed_weight = F.interpolate(pos_embed_weight,
size=input_shpae,
align_corners=False,
mode='bicubic')
cls_token_weight = cls_token_weight.unsqueeze(1)
pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
# pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
return pos_embed_weight.transpose(-2, -1)
def forward(self, x):
B, C, H, W = x.size()
res = self.connect(x)
x = x.reshape(B, C, -1) # NC(HW)
# x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(1+HW)
pos_embed = self.positional_embedding.unsqueeze(0)
pos_embed = self.resize_pos_embed(pos_embed, (H, W)) # NC(HW)
x = x + pos_embed.to(x.dtype) # NC(HW)
x = x.permute(2, 0, 1) # (HW)NC
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)
xt = x[0]
x = x.permute(1, 2, 0).reshape(B, -1, H, W)
x = x + res
x = F.relu(x, True)
return x, xt
"""
attenpool used in Clip (output: a tensor (b, dim) image encoding)
"""
class AttentionPool2d(nn.Module):
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.flatten(start_dim=2).permute(2, 0, 1) # NCHW -> (HW)NC
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)NC
x = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
x, _ = F.multi_head_attention_forward(
query=x[:1], 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.squeeze(0)
class ModifiedResNet(nn.Module):
"""
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
# the 3-layer stem
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)
# residual layers
self._inplanes = width # this is a *mutable* variable used during construction
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 # the ResNet feature dimension
self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim,
heads, output_dim)
# self.modifiedattnpool = ModifiedAttentionPool2d(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 = x.type(self.conv1.weight.dtype)
x = stem(x)
x = self.layer1(x)
x2 = self.layer2(x)
x3 = self.layer3(x2)
x4 = self.layer4(x3)
x5 = self.attnpool(x4)
# x4 = self.modifiedattnpool(x4)
return (x2, x3, x4), x5
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__()
# print(n_head)
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
res = self.attn(x, x, x, need_weights=False,
attn_mask=self.attn_mask)[0]
# print(res)
return res
def forward(self, x: torch.Tensor):
# a = self.attention(self.ln_1(x))
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):
super().__init__()
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 ViTTransformer(nn.Module):
def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
super().__init__()
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):
outputs = []
i = 1
for block in self.resblocks:
x = block(x)
if i > 7:
outputs.append(x)
i = i + 1
return outputs
class VisionTransformer(nn.Module):
def __init__(self, input_resolution: int, patch_size: int, width: int,
layers: int, heads: int, output_dim: int):
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 = ViTTransformer(width, layers, heads)
self.ln_post = LayerNorm(width)
self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
def forward(self, x: torch.Tensor):
# input: batch, 3, 224, 224
# batch, 1024, 16, 16
x = self.conv1(x) # shape = [*, width, grid, grid]
# batch, 1024, 256
x = x.reshape(x.shape[0], x.shape[1],
-1) # shape = [*, width, grid ** 2]
# batch, 256, 1024
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
# batch, 257, 1024
x = torch.cat([
self.class_embedding.to(x.dtype) + torch.zeros(
x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x
],
dim=1) # shape = [*, grid ** 2 + 1, width]
x = x + self.positional_embedding.to(x.dtype)
x = self.ln_pre(x)
# 257, batch, 1024
x = x.permute(1, 0, 2) # NLD -> LND
out = self.transformer(x)
# batch, 257, 1024
x1, x2 ,x3, x4 = out[0], out[1], out[2], out[3]
x1 = x1.permute(1, 0, 2)
x2 = x2.permute(1, 0, 2)
x3 = x3.permute(1, 0, 2)
x4 = x4.permute(1, 0, 2) # LND -> NLD
# 用于分类
x = self.ln_post(x4[:, 0, :])
#feature
# x_f = self.ln_post(x[:, 1:, :])
if self.proj is not None:
x = x @ self.proj
return (x1[:, 1:, :], x2[:, 1:, :], x3[:, 1:, :], x4[:, 1:, :]), x
class ModifiedVisionTransformer(nn.Module):
def __init__(self, input_resolution: int, patch_size: int, width: int,
layers: int, heads: int, output_dim: int):
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)
self.conv2 = nn.Conv2d(in_channels=3,
out_channels=width // 2,
kernel_size=patch_size // 2,
stride=patch_size // 2,
bias=False)
self.conv3 = nn.Conv2d(in_channels=3,
out_channels=width,
kernel_size=patch_size * 2,
stride=patch_size * 2,
bias=False)
self.conv_layers = [self.conv1, self.conv2]
scale = width**-0.5
self.class_embedding1 = nn.Parameter(scale * torch.randn(width))
self.class_embedding2 = nn.Parameter(scale * torch.randn(width // 2))
self.cls_layers = [self.class_embedding1, self.class_embedding2]
self.positional_embedding1 = nn.Parameter(scale * torch.randn(
(input_resolution // patch_size)**2 + 1, width))
self.positional_embedding2 = nn.Parameter(scale * torch.randn(
(input_resolution // (patch_size // 2)) ** 2 + 1, width // 2))
self.pos_layers = [self.positional_embedding1, self.positional_embedding2]
self.ln_pre1 = LayerNorm(width)
self.ln_pre2 = LayerNorm(width // 2)
self.pre_layers = [self.ln_pre1, self.ln_pre2]
self.transformer1 = Transformer(width, layers, heads)
self.transformer2 = Transformer(width // 2, layers, heads)
self.tran_layers = [self.transformer1, self.transformer2]
self.ln_post1 = LayerNorm(width)
self.ln_post2 = LayerNorm(width // 2)
self.post_layers = [self.ln_post1, self.ln_post2]
self.proj1 = nn.Parameter(scale * torch.randn(width, output_dim * 2))
self.proj2 = nn.Parameter(scale * torch.randn(width // 2, output_dim))
self.proj_layers = [self.proj1, self.proj2]
def forward(self, x: torch.Tensor):
# input: batch, 3, 224, 224
input = x
# batch, 1024, 16, 16
out = []
f = []
cl = []
for i in range(2):
x = self.conv_layers[i](input) # shape = [*, width, grid, grid]
b, c, w, h = x.shape
# batch, 1024, 256
x = x.reshape(x.shape[0], x.shape[1],
-1) # shape = [*, width, grid ** 2]
# batch, 256, 1024
x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
# batch, 257, 1024
x = torch.cat([
self.cls_layers[i].to(x.dtype) + torch.zeros(
x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x
],
dim=1) # shape = [*, grid ** 2 + 1, width]
x = x + self.pos_layers[i].to(x.dtype)
x = self.pre_layers[i](x)
# 257, batch, 1024
x = x.permute(1, 0, 2) # NLD -> LND
x, cls = self.tran_layers[i](x)
# batch, 257, 1024
x = x.permute(1, 0, 2) # LND -> NLD
# 用于分类
# x = self.ln_post(x[:, 0, :])
# feature
x = self.post_layers[i](x[:, 1:, :])
if self.proj_layers[i] is not None:
x = x @ self.proj_layers[i]
cls = [j @ self.proj_layers[i] for j in cls]
feat = x.permute(0,2,1).reshape(b, x.shape[2] , w, h)
out.append(x)
f.append(feat)
cl.append(cls)
return out, f, cl
"""
Long CLIP
"""
class LCLIP(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
context_length: int,
vocab_size: int,
transformer_width: int,
transformer_heads: int,
transformer_layers: int,
load_from_clip: bool
):
super().__init__()
self.context_length = 248
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 = VisionTransformer(
input_resolution=image_resolution,
patch_size=vision_patch_size,
width=vision_width,
layers=vision_layers,
heads=vision_heads,
output_dim=embed_dim
)
self.transformer = Transformer(
width=transformer_width,
layers=transformer_layers,
heads=transformer_heads,
attn_mask=self.build_attention_mask()
)
self.vocab_size = vocab_size
self.token_embedding = nn.Embedding(vocab_size, transformer_width)
# self.positional_embedding = nn.Parameter(torch.empty(248, transformer_width))
if load_from_clip == False:
self.positional_embedding = nn.Parameter(torch.empty(248, transformer_width))
self.positional_embedding_res = nn.Parameter(torch.empty(248, transformer_width))
else:
self.positional_embedding = nn.Parameter(torch.empty(248, transformer_width))
self.ln_final = LayerNorm(transformer_width)
self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
self.initialize_parameters()
self.mask1 = torch.zeros([248, 1])
self.mask1[:20, :] = 1
self.mask2 = torch.zeros([248, 1])
self.mask2[20:, :] = 1
def initialize_parameters(self):
nn.init.normal_(self.token_embedding.weight, std=0.02)
nn.init.normal_(self.positional_embedding, std=0.01)
if isinstance(self.visual, ModifiedResNet):
if self.visual.attnpool is not None:
std = self.visual.attnpool.c_proj.in_features ** -0.5
nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
for name, param in resnet_block.named_parameters():
if name.endswith("bn3.weight"):
nn.init.zeros_(param)
proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
attn_std = self.transformer.width ** -0.5
fc_std = (2 * self.transformer.width) ** -0.5
for block in self.transformer.resblocks:
nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
if self.text_projection is not None:
nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
def build_attention_mask(self):
# lazily create causal attention mask, with full attention between the vision tokens
# pytorch uses additive attention mask; fill with -inf
mask = torch.empty(self.context_length, self.context_length)
mask.fill_(float("-inf"))
mask.triu_(1) # zero out the lower diagonal
return mask
@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):
x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model]
# x = x + (self.positional_embedding.to(x.device) * self.mask1.to(x.device)).type(self.dtype).to(x.device) + (self.positional_embedding_res.to(x.device) * self.mask2.to(x.device)).type(self.dtype).to(x.device)
x = x + (self.positional_embedding.to(x.device) * self.mask1.to(x.device)).type(self.dtype).to(x.device)
x = x.permute(1, 0, 2) # NLD -> LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND -> NLD
x = self.ln_final(x).type(self.dtype)
# x.shape = [batch_size, n_ctx, transformer.width]
# take features from the eot embedding (eot_token is the highest number in each sequence)
x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
return x
def encode_text_full(self, text):
x = self.token_embedding(text).type(self.dtype) # [batch_size, n_ctx, d_model]
x = x + (self.positional_embedding.to(x.device) * self.mask1.to(x.device)).type(self.dtype).to(x.device) + (self.positional_embedding_res.to(x.device) * self.mask2.to(x.device)).type(self.dtype).to(x.device)
x = x.permute(1, 0, 2) # NLD -> LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND -> NLD
x = self.ln_final(x).type(self.dtype)
# x.shape = [batch_size, n_ctx, transformer.width]
# take features from the eot embedding (eot_token is the highest number in each sequence)
#x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
return x
def forward(self, image, text):
image_features = self.encode_image(image)
text_features, _ = self.encode_text(text)
# normalized features
image_features = image_features / image_features.norm(dim=1, keepdim=True)
text_features = text_features / text_features.norm(dim=1, keepdim=True)
# cosine similarity as logits
logit_scale = self.logit_scale.exp()
logits_per_image = logit_scale * image_features @ text_features.t()
logits_per_text = logits_per_image.t()
# shape = [global_batch_size, global_batch_size]
return logits_per_image, logits_per_text
"""
original CLIP
"""
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
context_length: int,
txt_length: int,
vocab_size: int,
transformer_width: int,
transformer_heads: int,
transformer_layers: int):
super().__init__()
self.context_length = context_length
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)
# self.fq_attnpool = AttentionPool2d(image_resolution // 32, vision_width* 32,
# vision_heads, embed_dim)
else:
vision_heads = vision_width // 64
self.visual = VisionTransformer(input_resolution=image_resolution,
patch_size=vision_patch_size,
width=vision_width,
layers=vision_layers,
heads=vision_heads,
output_dim=embed_dim)
self.transformer = Transformer(
width=transformer_width,
layers=transformer_layers,
heads=transformer_heads,
attn_mask=self.build_attention_mask(txt_length))
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.text_projection = nn.Parameter(
torch.empty(transformer_width, embed_dim))
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
self.token_embedding.requires_grad_ = False
self.initialize_parameters()
def initialize_parameters(self):
nn.init.normal_(self.token_embedding.weight, std=0.02)
nn.init.normal_(self.positional_embedding, std=0.01)
if isinstance(self.visual, ModifiedResNet):
if self.visual.attnpool is not None:
std = self.visual.attnpool.c_proj.in_features**-0.5
nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
for resnet_block in [
self.visual.layer1, self.visual.layer2, self.visual.layer3,
self.visual.layer4
]:
for name, param in resnet_block.named_parameters():
if name.endswith("bn3.weight"):
nn.init.zeros_(param)
proj_std = (self.transformer.width**-0.5) * (
(2 * self.transformer.layers)**-0.5)
attn_std = self.transformer.width**-0.5
fc_std = (2 * self.transformer.width)**-0.5
for block in self.transformer.resblocks:
nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
if self.text_projection is not None:
nn.init.normal_(self.text_projection,
std=self.transformer.width**-0.5)
def build_attention_mask(self, context_length):
# lazily create causal attention mask, with full attention between the vision tokens
# pytorch uses additive attention mask; fill with -inf
mask = torch.empty(context_length, context_length)
mask.fill_(float("-inf"))
mask.triu_(1) # zero out the lower diagonal
return mask
@property
def dtype(self):
return self.visual.conv1.weight.dtype
def encode_image(self, image):
return self.visual(image.type(self.dtype))
def encode_fq(self, image):
return self.fq_attnpool(image.type(self.dtype))
def encode_text(self, text):
a = self.token_embedding
x = self.token_embedding(text).type(
self.dtype) # [batch_size, n_ctx, d_model]
x = x + self.positional_embedding.type(self.dtype)[:x.size(1)]
# print(x.shape)
# print(x)
x = x.permute(1, 0, 2) # NLD -> LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND -> NLD
x = self.ln_final(x).type(self.dtype)
# print(text[0])
# x.shape = [batch_size, n_ctx, transformer.width]
# take features from the eot embedding (eot_token is the highest number in each sequence)
state = x[torch.arange(x.shape[0]),
text.argmax(dim=-1)] @ self.text_projection
# x = x @ self.text_projection
# state = x[torch.arange(x.shape[0]), text.argmax(dim=-1)]
return x, state
def forward(self, image, text):
image_features = self.encode_image(image)
text_features = self.encode_text(text)
# normalized features
image_features = image_features / image_features.norm(dim=-1,
keepdim=True)
text_features = text_features / text_features.norm(dim=-1,
keepdim=True)
# cosine similarity as logits
logit_scale = self.logit_scale.exp()
logits_per_image = logit_scale * image_features @ text_features.t()
logits_per_text = logits_per_image.t()
# shape = [global_batch_size, global_batch_size]
return logits_per_image, logits_per_text
"""
modified CLIP : without text encoder
"""
class zhCLIP(nn.Module):
def __init__(self,
embed_dim,
# vision
image_resolution: int,
vision_layers: Union[Tuple[int, int, int, int], int],
vision_width: int,
vision_patch_size: int):
super().__init__()
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)
self.fq_attnpool = AttentionPool2d(image_resolution // 32, vision_width* 32,
vision_heads, embed_dim)
else:
vision_heads = vision_width // 64
self.visual = ModifiedVisionTransformer(input_resolution=image_resolution,
patch_size=vision_patch_size,
width=vision_width,
layers=vision_layers,
heads=vision_heads,
output_dim=embed_dim)
self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
self.initialize_parameters()
def initialize_parameters(self):
if isinstance(self.visual, ModifiedResNet):
if self.visual.attnpool is not None:
std = self.visual.attnpool.c_proj.in_features**-0.5
nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
for resnet_block in [
self.visual.layer1, self.visual.layer2, self.visual.layer3,
self.visual.layer4
]:
for name, param in resnet_block.named_parameters():
if name.endswith("bn3.weight"):
nn.init.zeros_(param)
def build_attention_mask(self, context_length):
# lazily create causal attention mask, with full attention between the vision tokens
# pytorch uses additive attention mask; fill with -inf
mask = torch.empty(context_length, context_length)
mask.fill_(float("-inf"))
mask.triu_(1) # zero out the lower diagonal
return mask
@property
def dtype(self):
return self.visual.conv1.weight.dtype
def encode_image(self, image):
return self.visual(image.type(self.dtype))
def encode_fq(self, image):
return self.fq_attnpool(image.type(self.dtype))
def forward(self, image, text):
image_features = self.encode_image(image)
text_features = self.encode_text(text)
# normalized features
image_features = image_features / image_features.norm(dim=-1,
keepdim=True)
text_features = text_features / text_features.norm(dim=-1,
keepdim=True)
# cosine similarity as logits
logit_scale = self.logit_scale.exp()
logits_per_image = logit_scale * image_features @ text_features.t()
logits_per_text = logits_per_image.t()
# shape = [global_batch_size, global_batch_size]
return logits_per_image, logits_per_text
def convert_weights(model: nn.Module):
"""Convert applicable model parameters to fp16"""
def _convert_weights_to_fp16(l):
if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
l.weight.data = l.weight.data.half()
if l.bias is not None:
l.bias.data = l.bias.data.half()
if isinstance(l, nn.MultiheadAttention):
for attr in [
*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
"in_proj_bias", "bias_k", "bias_v"
]:
tensor = getattr(l, attr)
if tensor is not None:
tensor.data = tensor.data.half()
for name in ["text_projection", "proj"]:
if hasattr(l, name):
attr = getattr(l, name)
if attr is not None:
attr.data = attr.data.half()
model.apply(_convert_weights_to_fp16)
class PromptLearner(nn.Module):
def __init__(self, transformer_width, context_length, vocab_size,
transformer_layers, transformer_heads, bert_embed_dim):
super().__init__()
self.transformer_width = transformer_width
self.context_length = context_length
self.vocab_size = vocab_size
self.token_embedding = nn.Embedding(self.vocab_size, self.transformer_width)
self.transformer = Transformer(
width=transformer_width,
layers=transformer_layers,
heads=transformer_heads,
attn_mask=self.build_attention_mask()
)
self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
self.ln_final = LayerNorm(transformer_width)
self.text_projection = nn.Parameter(torch.empty(transformer_width, bert_embed_dim))
# self.load_from_openai_model(pretrained_model=clip_pretrain)
def build_attention_mask(self):
# lazily create causal attention mask, with full attention between the vision tokens
# pytorch uses additive attention mask; fill with -inf
mask = torch.empty(self.context_length, self.context_length)
mask.fill_(float("-inf"))
mask.triu_(1) # zero out the lower diagonal
return mask
def init_label_emb(self, labels_path):
label = open(labels_path, 'r').readlines()
# label81 = open(unseen_labels_path, 'r').readlines()
# label1006 = label925 + label81
self.name_lens = [len(_tokenizer.encode(name)) for name in label]
self.label_token = torch.zeros((len(self.name_lens), self.context_length), dtype=torch.long)
for i, c in enumerate(label):
self.label_token[i] = tokenize(f"There is a {c.strip()} in the scene")
self.label_emb = torch.zeros((len(self.name_lens), max(self.name_lens), self.transformer_width))
for i, embed in enumerate(self.token_embedding(self.label_token)):
self.label_emb[i][:self.name_lens[i]] = embed[4:4 + self.name_lens[i]].clone().detach()
# def load_from_openai_model(self, pretrained_model):
# state_dict = clip.load(pretrained_model, jit=False)[0].state_dict()
# load_dict = {}
# for k, v in state_dict.items():
# if not k.startswith("visual") and (
# k not in ["logit_scale", "input_resolution", "context_length", "vocab_size"]):
# load_dict[k] = v
# msg = self.load_state_dict(load_dict)
def load_label_emb(self, label=None):
self.name_lens = [len(_tokenizer.encode(name.split("\t")[-1])) for name in label]
self.label_token = torch.zeros((len(self.name_lens), self.context_length), dtype=torch.long).cuda()
for i, c in enumerate(label):
name = c.split("\t")[-1]
self.label_token[i] = tokenize(f"There is a {name.strip()} in the scene")
self.label_emb = torch.zeros((len(self.name_lens), max(self.name_lens), self.transformer_width)).cuda()
for i, embed in enumerate(self.token_embedding(self.label_token)):
self.label_emb[i][:self.name_lens[i]] = embed[4:4 + self.name_lens[i]].clone().detach()
def forward(self, device):
label_embeds = self.token_embedding(self.label_token.to(device))
for i in range(label_embeds.shape[0]):
label_embeds[i, 4:4 + self.name_lens[i], :] = self.label_emb[i][:self.name_lens[i]]
x = label_embeds + self.positional_embedding
x = x.permute(1, 0, 2) # NLD -> LND
x = self.transformer(x)
x = x.permute(1, 0, 2) # LND -> NLD
x = self.ln_final(x)
res = x[torch.arange(x.shape[0]), self.label_token.argmax(dim=-1)] @ self.text_projection
return res
def build_promptlearner(state_dict: dict):
embed_dim = state_dict["text_projection"].shape[1]
context_length = state_dict["positional_embedding"].shape[0]
vocab_size = state_dict["token_embedding.weight"].shape[0]
transformer_width = state_dict["ln_final.weight"].shape[0]
transformer_heads = transformer_width // 64
transformer_layers = len(
set(
k.split(".")[2] for k in state_dict
if k.startswith(f"transformer.resblocks")))
model = PromptLearner(transformer_width, context_length, vocab_size,
transformer_layers, transformer_heads, embed_dim)
# model = PromptLearner(embed_dim, vision_patch_size, context_length, txt_length, vocab_size,
# transformer_width, transformer_heads, transformer_layers)
load_dict = {}
for k, v in state_dict.items():
if not k.startswith("visual") and (
k not in ["logit_scale", "input_resolution", "context_length", "vocab_size"]):
load_dict[k] = v
convert_weights(model)
model.load_state_dict(load_dict, False)
return model
def build_model(state_dict: dict, txt_length: int):
vit = "visual.proj" in state_dict
if vit:
vision_width = state_dict["visual.conv1.weight"].shape[0]
vision_layers = len([
k for k in state_dict.keys()
if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")
])
vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
grid_size = round(
(state_dict["visual.positional_embedding"].shape[0] - 1)**0.5)
image_resolution = vision_patch_size * grid_size
else:
counts: list = [
len(
set(
k.split(".")[2] for k in state_dict
if k.startswith(f"visual.layer{b}")))
for b in [1, 2, 3, 4]
]
vision_layers = tuple(counts)
vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
output_width = round(
(state_dict["visual.attnpool.positional_embedding"].shape[0] -
1)**0.5)
vision_patch_size = None
assert output_width**2 + 1 == state_dict[
"visual.attnpool.positional_embedding"].shape[0]
image_resolution = output_width * 32
vision_heads = vision_width * 32 // 64
embed_dim = state_dict["text_projection"].shape[1]
# context_length = state_dict["positional_embedding"].shape[0]
context_length = txt_length
vocab_size = state_dict["token_embedding.weight"].shape[0]
transformer_width = state_dict["ln_final.weight"].shape[0]
transformer_heads = transformer_width // 64
transformer_layers = len(
set(
k.split(".")[2] for k in state_dict
if k.startswith(f"transformer.resblocks")))
model = CLIP(embed_dim, image_resolution, vision_layers, vision_width,
vision_patch_size, context_length, txt_length, vocab_size,
transformer_width, transformer_heads, transformer_layers)
for key in ["input_resolution", "context_length", "vocab_size", 'positional_embedding']:
if key in state_dict:
del state_dict[key]
convert_weights(model)
model.load_state_dict(state_dict, False)
return model.eval(), image_resolution, vision_heads, embed_dim, vision_width, vision_patch_size
def build_lclip_model(state_dict: dict, load_from_clip: bool):
vit = "visual.proj" in state_dict
if vit:
vision_width = state_dict["visual.conv1.weight"].shape[0]
vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
image_resolution = vision_patch_size * grid_size
else:
counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
vision_layers = tuple(counts)
vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
vision_patch_size = None
assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
image_resolution = output_width * 32
embed_dim = state_dict["text_projection"].shape[1]
# print(embed_dim)
context_length = state_dict["positional_embedding"].shape[0]
vocab_size = state_dict["token_embedding.weight"].shape[0]
transformer_width = state_dict["ln_final.weight"].shape[0]
transformer_heads = transformer_width // 64
transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
model = LCLIP(
embed_dim,
image_resolution, vision_layers, vision_width, vision_patch_size,
context_length, vocab_size, transformer_width, transformer_heads, transformer_layers, load_from_clip
)
for key in ["input_resolution", "context_length", "vocab_size"]:
if key in state_dict:
del state_dict[key]
convert_weights(model)
# model.load_state_dict(state_dict)
model.load_state_dict(state_dict, strict=False)
vision_heads = vision_width // 64
# print(vision_heads)
return model.eval(), image_resolution, vision_heads, embed_dim, vision_width, vision_patch_size
def build_modified_model(state_dict: dict, txt_length: int):
vit = "visual.proj" in state_dict
if vit:
vision_width = state_dict["visual.conv1.weight"].shape[0]
vision_layers = len([
k for k in state_dict.keys()
if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")
])
vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
grid_size = round(
(state_dict["visual.positional_embedding"].shape[0] - 1)**0.5)
image_resolution = vision_patch_size * grid_size
else:
counts: list = [
len(
set(
k.split(".")[2] for k in state_dict
if k.startswith(f"visual.layer{b}")))
for b in [1, 2, 3, 4]
]
vision_layers = tuple(counts)
vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
output_width = round(
(state_dict["visual.attnpool.positional_embedding"].shape[0] -
1)**0.5)
vision_patch_size = None
assert output_width**2 + 1 == state_dict[
"visual.attnpool.positional_embedding"].shape[0]
image_resolution = output_width * 32
embed_dim = state_dict["text_projection"].shape[1]
model = zhCLIP(embed_dim, image_resolution, vision_layers, vision_width,
vision_patch_size)
for key in ["input_resolution", "context_length", "vocab_size"]:
if key in state_dict:
del state_dict[key]
convert_weights(model)
model.load_state_dict(state_dict, False)
return model.eval()