AskAnything / intern_action.py
ynhe's picture
[Init]
5423e78
raw
history blame
19.1 kB
#!/usr/bin/env python
import os
from collections import OrderedDict
from timm.models.layers import DropPath
import torch
from torch import nn
from torch.nn import MultiheadAttention
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
MODEL_PATH = './'
_MODELS = {
"ViT-B/16": os.path.join(MODEL_PATH, "vit_b16.pth"),
"ViT-L/14": os.path.join(MODEL_PATH, "vit_l14.pth"),
"ViT-L/14_336": os.path.join(MODEL_PATH, "vit_l14_336.pth"),
}
class LayerNorm(nn.LayerNorm):
"""Subclass torch's LayerNorm to handle fp16."""
def forward(self, x):
orig_type = x.dtype
ret = super().forward(x.type(torch.float32))
return ret.type(orig_type)
class QuickGELU(nn.Module):
def forward(self, x):
return x * torch.sigmoid(1.702 * x)
class Local_MHRA(nn.Module):
def __init__(self, d_model, dw_reduction=1.5, pos_kernel_size=3):
super().__init__()
padding = pos_kernel_size // 2
re_d_model = int(d_model // dw_reduction)
self.pos_embed = nn.Sequential(
nn.BatchNorm3d(d_model),
nn.Conv3d(d_model, re_d_model, kernel_size=1, stride=1, padding=0),
nn.Conv3d(re_d_model, re_d_model, kernel_size=(pos_kernel_size, 1, 1), stride=(1, 1, 1), padding=(padding, 0, 0), groups=re_d_model),
nn.Conv3d(re_d_model, d_model, kernel_size=1, stride=1, padding=0),
)
# init zero
print('Init zero for Conv in pos_emb')
nn.init.constant_(self.pos_embed[3].weight, 0)
nn.init.constant_(self.pos_embed[3].bias, 0)
def forward(self, x):
return self.pos_embed(x)
class ResidualAttentionBlock(nn.Module):
def __init__(
self, d_model, n_head, attn_mask=None, drop_path=0.0,
dw_reduction=1.5, no_lmhra=False, double_lmhra=True
):
super().__init__()
self.n_head = n_head
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
print(f'Drop path rate: {drop_path}')
self.no_lmhra = no_lmhra
self.double_lmhra = double_lmhra
print(f'No L_MHRA: {no_lmhra}')
print(f'Double L_MHRA: {double_lmhra}')
if not no_lmhra:
self.lmhra1 = Local_MHRA(d_model, dw_reduction=dw_reduction)
if double_lmhra:
self.lmhra2 = Local_MHRA(d_model, dw_reduction=dw_reduction)
# spatial
self.attn = 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):
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, T=8, use_checkpoint=False):
# x: 1+HW, NT, C
if not self.no_lmhra:
# Local MHRA
tmp_x = x[1:, :, :]
L, NT, C = tmp_x.shape
N = NT // T
H = W = int(L ** 0.5)
tmp_x = tmp_x.view(H, W, N, T, C).permute(2, 4, 3, 0, 1).contiguous()
tmp_x = tmp_x + self.drop_path(self.lmhra1(tmp_x))
tmp_x = tmp_x.view(N, C, T, L).permute(3, 0, 2, 1).contiguous().view(L, NT, C)
x = torch.cat([x[:1, :, :], tmp_x], dim=0)
# MHSA
if use_checkpoint:
attn_out = checkpoint.checkpoint(self.attention, self.ln_1(x))
x = x + self.drop_path(attn_out)
else:
x = x + self.drop_path(self.attention(self.ln_1(x)))
# Local MHRA
if not self.no_lmhra and self.double_lmhra:
tmp_x = x[1:, :, :]
tmp_x = tmp_x.view(H, W, N, T, C).permute(2, 4, 3, 0, 1).contiguous()
tmp_x = tmp_x + self.drop_path(self.lmhra2(tmp_x))
tmp_x = tmp_x.view(N, C, T, L).permute(3, 0, 2, 1).contiguous().view(L, NT, C)
x = torch.cat([x[:1, :, :], tmp_x], dim=0)
# FFN
if use_checkpoint:
mlp_out = checkpoint.checkpoint(self.mlp, self.ln_2(x))
x = x + self.drop_path(mlp_out)
else:
x = x + self.drop_path(self.mlp(self.ln_2(x)))
return x
class Extractor(nn.Module):
def __init__(
self, d_model, n_head, attn_mask=None,
mlp_factor=4.0, dropout=0.0, drop_path=0.0,
):
super().__init__()
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
print(f'Drop path rate: {drop_path}')
self.attn = nn.MultiheadAttention(d_model, n_head)
self.ln_1 = nn.LayerNorm(d_model)
d_mlp = round(mlp_factor * d_model)
self.mlp = nn.Sequential(OrderedDict([
("c_fc", nn.Linear(d_model, d_mlp)),
("gelu", QuickGELU()),
("dropout", nn.Dropout(dropout)),
("c_proj", nn.Linear(d_mlp, d_model))
]))
self.ln_2 = nn.LayerNorm(d_model)
self.ln_3 = nn.LayerNorm(d_model)
self.attn_mask = attn_mask
# zero init
nn.init.xavier_uniform_(self.attn.in_proj_weight)
nn.init.constant_(self.attn.out_proj.weight, 0.)
nn.init.constant_(self.attn.out_proj.bias, 0.)
nn.init.xavier_uniform_(self.mlp[0].weight)
nn.init.constant_(self.mlp[-1].weight, 0.)
nn.init.constant_(self.mlp[-1].bias, 0.)
def attention(self, x, y):
d_model = self.ln_1.weight.size(0)
q = (x @ self.attn.in_proj_weight[:d_model].T) + self.attn.in_proj_bias[:d_model]
k = (y @ self.attn.in_proj_weight[d_model:-d_model].T) + self.attn.in_proj_bias[d_model:-d_model]
v = (y @ self.attn.in_proj_weight[-d_model:].T) + self.attn.in_proj_bias[-d_model:]
Tx, Ty, N = q.size(0), k.size(0), q.size(1)
q = q.view(Tx, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3)
k = k.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3)
v = v.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3)
aff = (q @ k.transpose(-2, -1) / (self.attn.head_dim ** 0.5))
aff = aff.softmax(dim=-1)
out = aff @ v
out = out.permute(2, 0, 1, 3).flatten(2)
out = self.attn.out_proj(out)
return out
def forward(self, x, y):
x = x + self.drop_path(self.attention(self.ln_1(x), self.ln_3(y)))
x = x + self.drop_path(self.mlp(self.ln_2(x)))
return x
class Transformer(nn.Module):
def __init__(
self, width, layers, heads, attn_mask=None, backbone_drop_path_rate=0.,
use_checkpoint=False, checkpoint_num=[0], t_size=8, dw_reduction=2,
no_lmhra=False, double_lmhra=True,
return_list=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
n_layers=12, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0.,
mlp_dropout=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
cls_dropout=0.5, num_classes=400,
):
super().__init__()
self.T = t_size
self.return_list = return_list
# backbone
b_dpr = [x.item() for x in torch.linspace(0, backbone_drop_path_rate, layers)]
self.resblocks = nn.ModuleList([
ResidualAttentionBlock(
width, heads, attn_mask,
drop_path=b_dpr[i],
dw_reduction=dw_reduction,
no_lmhra=no_lmhra,
double_lmhra=double_lmhra,
) for i in range(layers)
])
# checkpoint
self.use_checkpoint = use_checkpoint
self.checkpoint_num = checkpoint_num
self.n_layers = n_layers
print(f'Use checkpoint: {self.use_checkpoint}')
print(f'Checkpoint number: {self.checkpoint_num}')
# global block
assert n_layers == len(return_list)
if n_layers > 0:
self.temporal_cls_token = nn.Parameter(torch.zeros(1, 1, n_dim))
self.dpe = nn.ModuleList([
nn.Conv3d(n_dim, n_dim, kernel_size=3, stride=1, padding=1, bias=True, groups=n_dim)
for i in range(n_layers)
])
for m in self.dpe:
nn.init.constant_(m.bias, 0.)
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, n_layers)]
self.dec = nn.ModuleList([
Extractor(
n_dim, n_head, mlp_factor=mlp_factor,
dropout=mlp_dropout[i], drop_path=dpr[i],
) for i in range(n_layers)
])
self.balance = nn.Parameter(torch.zeros((n_dim)))
self.sigmoid = nn.Sigmoid()
# projection
self.proj = nn.Sequential(
nn.LayerNorm(n_dim),
nn.Dropout(cls_dropout),
nn.Linear(n_dim, num_classes),
)
def forward(self, x):
T_down = self.T
L, NT, C = x.shape
N = NT // T_down
H = W = int((L - 1) ** 0.5)
if self.n_layers > 0:
cls_token = self.temporal_cls_token.repeat(1, N, 1)
j = -1
for i, resblock in enumerate(self.resblocks):
if self.use_checkpoint and i < self.checkpoint_num[0]:
x = resblock(x, self.T, use_checkpoint=True)
else:
x = resblock(x, T_down)
if i in self.return_list:
j += 1
tmp_x = x.clone()
tmp_x = tmp_x.view(L, N, T_down, C)
# dpe
_, tmp_feats = tmp_x[:1], tmp_x[1:]
tmp_feats = tmp_feats.permute(1, 3, 2, 0).reshape(N, C, T_down, H, W)
tmp_feats = self.dpe[j](tmp_feats).view(N, C, T_down, L - 1).permute(3, 0, 2, 1).contiguous()
tmp_x[1:] = tmp_x[1:] + tmp_feats
# global block
tmp_x = tmp_x.permute(2, 0, 1, 3).flatten(0, 1) # T * L, N, C
cls_token = self.dec[j](cls_token, tmp_x)
if self.n_layers > 0:
weight = self.sigmoid(self.balance)
residual = x.view(L, N, T_down, C)[0].mean(1) # L, N, T, C
return self.proj((1 - weight) * cls_token[0, :, :] + weight * residual)
else:
residual = x.view(L, N, T_down, C)[0].mean(1) # L, N, T, C
return self.proj(residual)
class VisionTransformer(nn.Module):
def __init__(
self,
# backbone
input_resolution, patch_size, width, layers, heads, output_dim, backbone_drop_path_rate=0.,
use_checkpoint=False, checkpoint_num=[0], t_size=8, kernel_size=3, dw_reduction=1.5,
temporal_downsample=True,
no_lmhra=-False, double_lmhra=True,
# global block
return_list=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
n_layers=12, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0.,
mlp_dropout=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
cls_dropout=0.5, num_classes=400,
):
super().__init__()
self.input_resolution = input_resolution
self.output_dim = output_dim
padding = (kernel_size - 1) // 2
if temporal_downsample:
self.conv1 = nn.Conv3d(3, width, (kernel_size, patch_size, patch_size), (2, patch_size, patch_size), (padding, 0, 0), bias=False)
t_size = t_size // 2
else:
self.conv1 = nn.Conv3d(3, width, (1, patch_size, patch_size), (1, patch_size, patch_size), (0, 0, 0), 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, dw_reduction=dw_reduction,
backbone_drop_path_rate=backbone_drop_path_rate,
use_checkpoint=use_checkpoint, checkpoint_num=checkpoint_num, t_size=t_size,
no_lmhra=no_lmhra, double_lmhra=double_lmhra,
return_list=return_list, n_layers=n_layers, n_dim=n_dim, n_head=n_head,
mlp_factor=mlp_factor, drop_path_rate=drop_path_rate, mlp_dropout=mlp_dropout,
cls_dropout=cls_dropout, num_classes=num_classes,
)
def forward(self, x):
x = self.conv1(x) # shape = [*, width, grid, grid]
N, C, T, H, W = x.shape
x = x.permute(0, 2, 3, 4, 1).reshape(N * T, H * W, C)
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)
x = x.permute(1, 0, 2) # NLD -> LND
out = self.transformer(x)
return out
def inflate_weight(weight_2d, time_dim, center=True):
print(f'Init center: {center}')
if center:
weight_3d = torch.zeros(*weight_2d.shape)
weight_3d = weight_3d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
middle_idx = time_dim // 2
weight_3d[:, :, middle_idx, :, :] = weight_2d
else:
weight_3d = weight_2d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1)
weight_3d = weight_3d / time_dim
return weight_3d
def load_state_dict(model, state_dict):
state_dict_3d = model.state_dict()
for k in state_dict.keys():
if state_dict[k].shape != state_dict_3d[k].shape:
if len(state_dict_3d[k].shape) <= 2:
print(f'Ignore: {k}')
continue
print(f'Inflate: {k}, {state_dict[k].shape} => {state_dict_3d[k].shape}')
time_dim = state_dict_3d[k].shape[2]
state_dict[k] = inflate_weight(state_dict[k], time_dim)
model.load_state_dict(state_dict, strict=False)
def intern_action_b16(
pretrained=True, use_checkpoint=False, checkpoint_num=[0],
t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0.,
temporal_downsample=True,
no_lmhra=False, double_lmhra=True,
return_list=[8, 9, 10, 11],
n_layers=4, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0.,
mlp_dropout=[0.5, 0.5, 0.5, 0.5],
cls_dropout=0.5, num_classes=400,
):
model = VisionTransformer(
input_resolution=224,
patch_size=16,
width=768,
layers=12,
heads=12,
output_dim=512,
use_checkpoint=use_checkpoint,
checkpoint_num=checkpoint_num,
t_size=t_size,
dw_reduction=dw_reduction,
backbone_drop_path_rate=backbone_drop_path_rate,
temporal_downsample=temporal_downsample,
no_lmhra=no_lmhra,
double_lmhra=double_lmhra,
return_list=return_list,
n_layers=n_layers,
n_dim=n_dim,
n_head=n_head,
mlp_factor=mlp_factor,
drop_path_rate=drop_path_rate,
mlp_dropout=mlp_dropout,
cls_dropout=cls_dropout,
num_classes=num_classes,
)
if pretrained:
print('load pretrained weights')
state_dict = torch.load(_MODELS["ViT-B/16"], map_location='cpu')
load_state_dict(model, state_dict)
return model.eval()
def intern_action_l14(
pretrained=True, use_checkpoint=False, checkpoint_num=[0],
t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0.,
temporal_downsample=True,
no_lmhra=False, double_lmhra=True,
return_list=[20, 21, 22, 23],
n_layers=4, n_dim=1024, n_head=16, mlp_factor=4.0, drop_path_rate=0.,
mlp_dropout=[0.5, 0.5, 0.5, 0.5],
cls_dropout=0.5, num_classes=400,
):
model = VisionTransformer(
input_resolution=224,
patch_size=14,
width=1024,
layers=24,
heads=16,
output_dim=768,
use_checkpoint=use_checkpoint,
checkpoint_num=checkpoint_num,
t_size=t_size,
dw_reduction=dw_reduction,
backbone_drop_path_rate=backbone_drop_path_rate,
temporal_downsample=temporal_downsample,
no_lmhra=no_lmhra,
double_lmhra=double_lmhra,
return_list=return_list,
n_layers=n_layers,
n_dim=n_dim,
n_head=n_head,
mlp_factor=mlp_factor,
drop_path_rate=drop_path_rate,
mlp_dropout=mlp_dropout,
cls_dropout=cls_dropout,
num_classes=num_classes,
)
if pretrained:
print('load pretrained weights')
state_dict = torch.load(_MODELS["ViT-L/14"], map_location='cpu')
load_state_dict(model, state_dict)
return model.eval()
def intern_action_l14_336(
pretrained=True, use_checkpoint=False, checkpoint_num=[0],
t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0.,
no_temporal_downsample=True,
no_lmhra=False, double_lmhra=True,
return_list=[20, 21, 22, 23],
n_layers=4, n_dim=1024, n_head=16, mlp_factor=4.0, drop_path_rate=0.,
mlp_dropout=[0.5, 0.5, 0.5, 0.5],
cls_dropout=0.5, num_classes=400,
):
model = VisionTransformer(
input_resolution=336,
patch_size=14,
width=1024,
layers=24,
heads=16,
output_dim=768,
use_checkpoint=use_checkpoint,
checkpoint_num=checkpoint_num,
t_size=t_size,
dw_reduction=dw_reduction,
backbone_drop_path_rate=backbone_drop_path_rate,
no_temporal_downsample=no_temporal_downsample,
no_lmhra=no_lmhra,
double_lmhra=double_lmhra,
return_list=return_list,
n_layers=n_layers,
n_dim=n_dim,
n_head=n_head,
mlp_factor=mlp_factor,
drop_path_rate=drop_path_rate,
mlp_dropout=mlp_dropout,
cls_dropout=cls_dropout,
num_classes=num_classes,
)
if pretrained:
print('load pretrained weights')
state_dict = torch.load(_MODELS["ViT-L/14_336"], map_location='cpu')
load_state_dict(model, state_dict)
return model.eval()
if __name__ == '__main__':
import time
from fvcore.nn import FlopCountAnalysis
from fvcore.nn import flop_count_table
import numpy as np
seed = 4217
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
num_frames = 16
model = intern_action_l14(
pretrained=False,
t_size=num_frames, backbone_drop_path_rate=0., drop_path_rate=0.,
dw_reduction=1.5,
no_lmhra=False,
temporal_downsample=True,
return_list=[8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
mlp_dropout=[0.5]*16,
n_layers=16
)
print(model)
flops = FlopCountAnalysis(model, torch.rand(1, 3, num_frames, 224, 224))
s = time.time()
print(flop_count_table(flops, max_depth=1))
print(time.time()-s)