253 lines
8.3 KiB
Python
253 lines
8.3 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) 2021 OpenAI
|
|
#
|
|
# This source code is licensed under the MIT license which can be found at
|
|
# https://github.com/openai/CLIP/blob/main/LICENSE
|
|
from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from fairseq.modules import LayerNorm
|
|
from torch import nn
|
|
|
|
from .utils.utils import DropPath
|
|
|
|
__all__ = [
|
|
'vit_base',
|
|
'vit_large',
|
|
'vit_large_336',
|
|
'vit_huge',
|
|
]
|
|
|
|
|
|
class QuickGELU(nn.Module):
|
|
r"""
|
|
An activation function module.
|
|
"""
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return x * torch.sigmoid(1.702 * x)
|
|
|
|
|
|
class ResidualAttentionBlock(nn.Module):
|
|
r"""
|
|
A residual attention block module.
|
|
|
|
step 1. Calculate the self attention in input with layer normalization.
|
|
step 2. Add input to the result of self attention's result as I.
|
|
step 3. Calculate the mlp of input I with layer normalization.
|
|
step 4. Add I to the result of mlp.
|
|
"""
|
|
|
|
def __init__(self,
|
|
d_model: int,
|
|
n_head: int,
|
|
attn_mask: torch.Tensor = None,
|
|
drop_path_rate=0.0):
|
|
r"""
|
|
Args:
|
|
d_model (`int`): The embedding dimensions.
|
|
n_head (`int`): The number of heads in self attention block.
|
|
attn_mask (`Tensor`, **optional**, default to None):
|
|
Attention mask using in self attention.
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
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
|
|
self.drop_path = DropPath(drop_path_rate)
|
|
|
|
def attention(self, x: torch.Tensor):
|
|
r"""
|
|
A wrapper of self attention .
|
|
"""
|
|
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.drop_path(self.attention(self.ln_1(x)))
|
|
x = x + self.drop_path(self.mlp(self.ln_2(x)))
|
|
return x
|
|
|
|
|
|
class Transformer(nn.Module):
|
|
r"""
|
|
A transformer module using in `VisionTransformer`.
|
|
|
|
Execute a sequential of `ResidualAttentionBlock`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
width: int,
|
|
layers: int,
|
|
heads: int,
|
|
attn_mask: torch.Tensor = None,
|
|
drop_path_rate: float = 0.0,
|
|
):
|
|
r"""
|
|
Args:
|
|
width (`int`): The width of input image.
|
|
layers (`int`): The number of `ResidualAttentionBlock` layers.
|
|
heads (int): The number of self attention heads.
|
|
attn_mask (`Tensor`, **optional**, default to None):
|
|
Attention mask using in self attention.
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
super().__init__()
|
|
self.width = width
|
|
self.layers = layers
|
|
self.resblocks = nn.Sequential(*[
|
|
ResidualAttentionBlock(width, heads, attn_mask, drop_path_rate)
|
|
for _ in range(layers)
|
|
])
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return self.resblocks(x)
|
|
|
|
|
|
class VisionTransformer(nn.Module):
|
|
r"""
|
|
Vision transformer module.
|
|
|
|
step 1. Using conv2d to get the image embedding.
|
|
step 2. If the resolution of input image doesn't equal to the initialized one
|
|
do `bilinear` interpolate to get new patch position embedding.
|
|
step 3. Add position embedding to image embedding to generate final image representation.
|
|
step 4. Do `Transformer` to the image representation.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_resolution: int,
|
|
patch_size: int,
|
|
width: int,
|
|
layers: int,
|
|
heads: int,
|
|
drop_path_rate: float = 0.0,
|
|
):
|
|
r"""
|
|
Args:
|
|
input_resolution (`int`): The resolution of input image.
|
|
patch_size (`int`): The resolution of each patch image.
|
|
width (`int`): The dimension of each patch image.
|
|
layers (`int`): The number of `ResidualAttentionBlock` in `Transformer`.
|
|
heads (`int`): The number of heads in self attention block.
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
super().__init__()
|
|
self.input_resolution = input_resolution
|
|
self.patch_size = patch_size
|
|
self.conv1 = nn.Conv2d(
|
|
in_channels=3,
|
|
out_channels=width,
|
|
kernel_size=patch_size,
|
|
stride=patch_size,
|
|
bias=False,
|
|
)
|
|
|
|
scale = width**-0.5
|
|
self.width = 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, drop_path_rate=drop_path_rate)
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
resolution = x.shape[-2]
|
|
height, width = x.shape[-2] // self.patch_size, x.shape[
|
|
-1] // self.patch_size
|
|
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]
|
|
|
|
if resolution != self.input_resolution:
|
|
old_pe = self.positional_embedding[1:]
|
|
patch_num = self.input_resolution // self.patch_size
|
|
old_pe = old_pe.reshape(1, patch_num, patch_num,
|
|
-1).permute(0, 3, 1, 2)
|
|
new_pe = F.interpolate(
|
|
old_pe, size=(height, width), mode='bilinear')
|
|
new_pe = new_pe.permute(0, 2, 3, 1).reshape(height * width, -1)
|
|
x = x + new_pe.to(x.dtype)
|
|
else:
|
|
x = x + self.positional_embedding[1:].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
|
|
|
|
bz, seq, hidden = x.shape
|
|
x = x.transpose(1, 2).reshape(bz, hidden, height, width)
|
|
|
|
return x
|
|
|
|
|
|
def vit_base(drop_path_rate: float = 0.0):
|
|
r"""
|
|
An instance of base vision transformer model.
|
|
|
|
Args:
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
return VisionTransformer(224, 16, 768, 9, 12, drop_path_rate)
|
|
|
|
|
|
def vit_large(drop_path_rate: float = 0.0):
|
|
r"""
|
|
An instance of large vision transformer model.
|
|
|
|
Args:
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
return VisionTransformer(224, 14, 1024, 18, 16, drop_path_rate)
|
|
|
|
|
|
def vit_large_336(drop_path_rate: float = 0.0):
|
|
r"""
|
|
An instance of large vision transformer model with 336 as input image width .
|
|
|
|
Args:
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
return VisionTransformer(336, 14, 1024, 18, 16, drop_path_rate)
|
|
|
|
|
|
def vit_huge(drop_path_rate: float = 0.0):
|
|
r"""
|
|
An instance of huge vision transformer model.
|
|
|
|
Args:
|
|
drop_path_rate (`float`, **optional**, default to 0.0):
|
|
Drop path rate. See more details about drop path from
|
|
https://arxiv.org/pdf/1605.07648v4.pdf.
|
|
"""
|
|
return VisionTransformer(224, 14, 1280, 24, 16, drop_path_rate)
|