WishArdently/InternVideo2Stage2-VisionEncoder
1128
1import math2import torch3import torch.nn.functional as F4from timm.models.layers import DropPath, to_2tuple, trunc_normal_5from torch import nn6 7import torch.utils.checkpoint as checkpoint8from functools import partial9from einops import rearrange10 11from .pos_embed import get_3d_sincos_pos_embed, get_2d_sincos_pos_embed, get_1d_sincos_pos_embed, interpolate_pos_embed_internvideo212from .flash_attention_class import FlashAttention13 14from transformers.utils import logging as error_logging15 16# Set up logging17error_logging.set_verbosity_error()18 19try:20 from flash_attn.modules.mlp import Mlp as FusedMLP21except:22 pass23 24try:25 from flash_attn.ops.rms_norm import DropoutAddRMSNorm26except:27 pass28 29 30class CrossAttention(nn.Module):31 def __init__(32 self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,33 proj_drop=0., attn_head_dim=None, out_dim=None):34 super().__init__()35 if out_dim is None:36 out_dim = dim37 self.num_heads = num_heads38 head_dim = dim // num_heads39 if attn_head_dim is not None:40 head_dim = attn_head_dim41 all_head_dim = head_dim * self.num_heads42 self.scale = qk_scale or head_dim ** -0.543 assert all_head_dim == dim44 45 self.q = nn.Linear(dim, all_head_dim, bias=False)46 self.k = nn.Linear(dim, all_head_dim, bias=False)47 self.v = nn.Linear(dim, all_head_dim, bias=False)48 49 if qkv_bias:50 self.q_bias = nn.Parameter(torch.zeros(all_head_dim))51 self.k_bias = nn.Parameter(torch.zeros(all_head_dim))52 self.v_bias = nn.Parameter(torch.zeros(all_head_dim))53 else:54 self.q_bias = None55 self.k_bias = None56 self.v_bias = None57 58 self.attn_drop = nn.Dropout(attn_drop)59 self.proj = nn.Linear(all_head_dim, out_dim)60 self.proj_drop = nn.Dropout(proj_drop)61 62 def forward(self, x, k=None, v=None):63 B, N, C = x.shape64 N_k = k.shape[1]65 N_v = v.shape[1]66 67 q_bias, k_bias, v_bias = None, None, None68 if self.q_bias is not None:69 q_bias = self.q_bias70 k_bias = self.k_bias71 v_bias = self.v_bias72 73 q = F.linear(input=x, weight=self.q.weight, bias=q_bias)74 q = q.reshape(B, N, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0) # (B, N_head, N_q, dim)75 76 k = F.linear(input=k, weight=self.k.weight, bias=k_bias)77 k = k.reshape(B, N_k, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)78 79 v = F.linear(input=v, weight=self.v.weight, bias=v_bias)80 v = v.reshape(B, N_v, 1, self.num_heads, -1).permute(2, 0, 3, 1, 4).squeeze(0)81 82 q = q * self.scale83 attn = (q @ k.transpose(-2, -1)) # (B, N_head, N_q, N_k)84 85 attn = attn.softmax(dim=-1)86 attn = self.attn_drop(attn)87 88 x = (attn @ v).transpose(1, 2).reshape(B, N, -1)89 x = self.proj(x)90 x = self.proj_drop(x)91 92 return x93 94 95class AttentiveBlock(nn.Module):96 97 def __init__(self, dim, num_heads, qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,98 drop_path=0., norm_layer=nn.LayerNorm, attn_head_dim=None, out_dim=None):99 super().__init__()100 101 self.norm1_q = norm_layer(dim)102 self.norm1_k = norm_layer(dim)103 self.norm1_v = norm_layer(dim)104 self.cross_attn = CrossAttention(105 dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,106 proj_drop=drop, attn_head_dim=attn_head_dim, out_dim=out_dim)107 108 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()109 110 def forward(self, x_q, x_kv, pos_q, pos_k, bool_masked_pos, rel_pos_bias=None):111 x_q = self.norm1_q(x_q + pos_q)112 x_k = self.norm1_k(x_kv + pos_k)113 x_v = self.norm1_v(x_kv)114 x = self.cross_attn(x_q, k=x_k, v=x_v)115 116 return x117 118 119class AttentionPoolingBlock(AttentiveBlock):120 121 def forward(self, x):122 # x_q = x.mean(1, keepdim=True)123 x_q = x124 x_kv, pos_q, pos_k = x, 0, 0125 x = super().forward(x_q, x_kv, pos_q, pos_k, bool_masked_pos=None, rel_pos_bias=None)126 x = x.squeeze(1)127 return x128 129 130class RMSNorm(nn.Module):131 def __init__(self, hidden_size, eps=1e-6):132 super().__init__()133 self.weight = nn.Parameter(torch.ones(hidden_size))134 self.variance_epsilon = eps135 136 def forward(self, hidden_states):137 input_dtype = hidden_states.dtype138 hidden_states = hidden_states.to(torch.float32)139 variance = hidden_states.pow(2).mean(-1, keepdim=True)140 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)141 return self.weight * hidden_states.to(input_dtype)142 143 144class LayerScale(nn.Module):145 def __init__(self, dim, init_values=1e-5, inplace=False, force_fp32=False):146 super().__init__()147 self.inplace = inplace148 self.gamma = nn.Parameter(init_values * torch.ones(dim))149 self.force_fp32 = force_fp32150 151 @torch.cuda.amp.autocast(enabled=False)152 def forward(self, x):153 if self.force_fp32:154 output_type = x.dtype155 out = x.float().mul_(self.gamma.float()) if self.inplace else x.float() * self.gamma.float()156 return out.to(dtype=output_type)157 else:158 out = x.mul_(self.gamma) if self.inplace else x * self.gamma159 return out160 161 162class Attention(nn.Module):163 def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., use_flash_attn=False,164 causal=False, norm_layer=nn.LayerNorm, qk_normalization=False, use_fused_rmsnorm=False):165 super().__init__()166 assert dim % num_heads == 0, 'dim should be divisible by num_heads'167 self.num_heads = num_heads168 head_dim = dim // num_heads169 self.scale = head_dim ** -0.5170 171 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)172 self.attn_drop = nn.Dropout(attn_drop)173 self.proj = nn.Linear(dim, dim)174 self.proj_drop = nn.Dropout(proj_drop)175 176 self.use_flash_attn = use_flash_attn177 if use_flash_attn:178 self.causal = causal179 self.inner_attn = FlashAttention(attention_dropout=attn_drop)180 181 self.qk_normalization = qk_normalization182 self.q_norm = norm_layer(dim) if qk_normalization else nn.Identity()183 self.k_norm = norm_layer(dim) if qk_normalization else nn.Identity()184 self.use_fused_rmsnorm = use_fused_rmsnorm185 186 def _naive_attn(self, x):187 B, N, C = x.shape188 # print(x.shape, torch.cuda.memory_allocated(), torch.cuda.memory_allocated())189 qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)190 q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple)191 192 if self.qk_normalization:193 B_, H_, N_, D_ = q.shape194 q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)195 k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)196 197 attn = ((q * self.scale) @ k.transpose(-2, -1))198 # attn = attn - attn.max(-1)[0].unsqueeze(-1) # in case of overflow for fp16199 attn = attn.softmax(dim=-1)200 attn = self.attn_drop(attn)201 # print(torch.cuda.memory_allocated(), torch.cuda.memory_allocated())202 x = (attn @ v).transpose(1, 2).reshape(B, N, C)203 x = self.proj(x)204 x = self.proj_drop(x)205 return x206 207 def _flash_attn(self, x, key_padding_mask=None, need_weights=False):208 209 qkv = self.qkv(x)210 qkv = rearrange(qkv, "b s (three h d) -> b s three h d", three=3, h=self.num_heads)211 212 if self.qk_normalization:213 q, k, v = qkv.unbind(2)214 if self.use_fused_rmsnorm:215 q = self.q_norm(q.flatten(-2, -1))[0].view(q.shape)216 k = self.k_norm(k.flatten(-2, -1))[0].view(k.shape)217 else:218 q = self.q_norm(q.flatten(-2, -1)).view(q.shape)219 k = self.k_norm(k.flatten(-2, -1)).view(k.shape)220 qkv = torch.stack([q, k, v], dim=2)221 222 context, _ = self.inner_attn(223 qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=self.causal224 )225 outs = self.proj(rearrange(context, "b s h d -> b s (h d)"))226 outs = self.proj_drop(outs)227 return outs228 229 def forward(self, x):230 x = self._naive_attn(x) if not self.use_flash_attn else self._flash_attn(x)231 return x232 233 234class Mlp(nn.Module):235 """ MLP as used in Vision Transformer, MLP-Mixer and related networks236 """237 238 def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU,239 bias=True, drop=0.):240 super().__init__()241 out_features = out_features or in_features242 hidden_features = hidden_features or in_features243 bias = to_2tuple(bias)244 drop_probs = to_2tuple(drop)245 246 self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])247 self.act = act_layer()248 self.drop1 = nn.Dropout(drop_probs[0])249 self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])250 self.drop2 = nn.Dropout(drop_probs[1])251 252 def forward(self, x):253 x = self.fc1(x)254 x = self.act(x)255 x = self.drop1(x)256 x = self.fc2(x)257 x = self.drop2(x)258 return x259 260 261class Block(nn.Module):262 263 def __init__(264 self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., init_values=None,265 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_flash_attn=False, use_fused_mlp=False,266 fused_mlp_heuristic=1, with_cp=False, qk_normalization=False, layerscale_no_force_fp32=False,267 use_fused_rmsnorm=False):268 super().__init__()269 270 self.norm1 = norm_layer(dim)271 self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,272 use_flash_attn=use_flash_attn, causal=False, norm_layer=norm_layer,273 qk_normalization=qk_normalization,274 use_fused_rmsnorm=use_fused_rmsnorm)275 self.ls1 = LayerScale(dim, init_values=init_values,276 force_fp32=(not layerscale_no_force_fp32)) if init_values else nn.Identity()277 # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here278 self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()279 280 self.norm2 = norm_layer(dim)281 mlp_hidden_dim = int(dim * mlp_ratio)282 if use_fused_mlp:283 # self.mlp = FusedMLP(in_features=dim, hidden_features=mlp_hidden_dim, heuristic=fused_mlp_heuristic)284 self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)285 else:286 self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)287 self.ls2 = LayerScale(dim, init_values=init_values,288 force_fp32=(not layerscale_no_force_fp32)) if init_values else nn.Identity()289 self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()290 291 self.with_cp = with_cp292 self.use_fused_rmsnorm = use_fused_rmsnorm293 294 def forward(self, x, residual=None):295 296 def _inner_forward(x, residual=None):297 if self.use_fused_rmsnorm:298 x, residual = self.norm1(x, residual)299 x = self.drop_path1(self.ls1(self.attn(x)))300 x, residual = self.norm2(x, residual)301 x = self.drop_path2(self.ls2(self.mlp(x)))302 return x, residual303 else:304 assert residual is None305 x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))306 x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))307 return x308 309 if self.with_cp:310 # print(f"\033[31m use_checkpoint [0m")311 return checkpoint.checkpoint(_inner_forward, x, residual)312 else:313 return _inner_forward(x, residual=residual)314 315 316class PatchEmbed(nn.Module):317 """ 3D Image to Patch Embedding318 """319 320 def __init__(321 self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, 322 num_frames=8, tubelet_size=1, norm_layer=None323 ):324 super().__init__()325 img_size = to_2tuple(img_size)326 patch_size = to_2tuple(patch_size)327 self.img_size = img_size328 self.patch_size = patch_size329 self.grid_size = (330 num_frames // tubelet_size, 331 img_size[0] // patch_size[0], 332 img_size[1] // patch_size[1]333 ) # (T, H, W)334 self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2]335 self.num_img_patches = self.grid_size[1] * self.grid_size[2]336 337 self.proj = nn.Conv3d(338 in_channels=in_chans, out_channels=embed_dim, 339 kernel_size=(tubelet_size, patch_size[0], patch_size[1]), 340 stride=(tubelet_size, patch_size[0], patch_size[1])341 )342 self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()343 344 def forward(self, x):345 x = self.proj(x)346 x = x.flatten(3).permute(0, 2, 3, 1) # B x C x T x HW => B x T x HW x C347 x = self.norm(x)348 return x349 350 351class Linear_Decoder(nn.Module):352 def __init__(self, in_channels=1408, out_channels=3200, 353 norm_layer=nn.LayerNorm, clip_norm_type='l2'):354 super().__init__()355 self.clip_norm_type = clip_norm_type356 # logger.info(f'Normalization Type: {clip_norm_type}')357 358 self.head = nn.Linear(in_channels, out_channels)359 self.norm = norm_layer(out_channels)360 361 self.apply(self._init_weights)362 363 def _init_weights(self, m):364 if isinstance(m, nn.Linear):365 nn.init.xavier_uniform_(m.weight)366 if isinstance(m, nn.Linear) and m.bias is not None:367 nn.init.constant_(m.bias, 0)368 elif isinstance(m, nn.LayerNorm):369 nn.init.constant_(m.bias, 0)370 nn.init.constant_(m.weight, 1.0)371 372 def forward(self, x):373 x = self.norm(self.head(x))374 375 if self.clip_norm_type == 'l2':376 x = x / x.norm(dim=-1, keepdim=True)377 elif self.clip_norm_type == 'none':378 pass379 else:380 raise NotImplementedError381 382 return x383 384 385class PretrainInternVideo2(nn.Module):386 def __init__(387 self,388 in_chans: int = 3,389 patch_size: int = 14,390 img_size: int = 224,391 qkv_bias: bool = False,392 drop_path_rate: float = 0.25,393 embed_dim: int = 1408,394 num_heads: int = 16,395 mlp_ratio: float = 48/11,396 init_values: float = 1e-5,397 qk_normalization: bool = True,398 depth: int = 40,399 use_flash_attn: bool = True,400 use_fused_rmsnorm: bool = True,401 use_fused_mlp: bool = True,402 fused_mlp_heuristic: int = 1,403 attn_pool_num_heads: int = 16,404 clip_embed_dim: int = 768,405 layerscale_no_force_fp32: bool = False,406 num_frames: int = 8,407 tubelet_size: int = 1,408 sep_pos_embed: bool = False,409 sep_image_video_pos_embed: bool = False,410 use_checkpoint: bool = False,411 checkpoint_num: int = 0,412 # for unmasked teacher413 clip_teacher_embed_dim: int = 3200,414 clip_teacher_final_dim: int = 768, # if 0, not distill final features415 clip_norm_type: str = 'l2',416 clip_return_layer: int = 1,417 clip_student_return_interval: int = 1,418 ):419 super().__init__()420 421 self.num_frames = num_frames422 # print(f'num_frames: {num_frames}')423 self.tubelet_size = tubelet_size424 assert use_flash_attn == use_fused_rmsnorm == use_fused_mlp, 'use_flash_attn, use_fused_rmsnorm and use_fused_mlp should be consistent'425 426 self.use_flash_attn = use_flash_attn427 self.embed_dim = embed_dim428 429 self.depth = depth430 self.clip_norm_type = clip_norm_type431 self.return_index = []432 for i in range(clip_return_layer):433 self.return_index.append(depth - int(i * clip_student_return_interval) - 1)434 # logger.info(f'Normalization Type: {clip_norm_type}')435 # logger.info(f'Strudent Return Index: {self.return_index}')436 437 if use_fused_rmsnorm:438 norm_layer_for_blocks = partial(DropoutAddRMSNorm, eps=1e-6, prenorm=True)439 else:440 norm_layer_for_blocks = partial(RMSNorm, eps=1e-6)441 self.norm_layer_for_blocks = norm_layer_for_blocks442 self.patch_embed = PatchEmbed(443 img_size, patch_size, in_chans, embed_dim,444 num_frames=num_frames, tubelet_size=tubelet_size,445 )446 num_patches = self.patch_embed.num_patches447 num_img_patches = self.patch_embed.num_img_patches448 449 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))450 451 # stolen from https://github.com/facebookresearch/mae_st/blob/dc072aaaf640d06892e23a33b42223a994efe272/models_vit.py#L65-L73C17452 self.sep_pos_embed = sep_pos_embed453 self.sep_image_video_pos_embed = sep_image_video_pos_embed454 if sep_pos_embed:455 raise NotImplementedError456 else:457 if sep_image_video_pos_embed:458 # logger.info("Use joint position embedding, for image and video we use different pos_embed.")459 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))460 self.img_pos_embed = nn.Parameter(torch.zeros(1, num_img_patches + 1, embed_dim))461 # for CLIP decoder462 self.clip_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))463 self.clip_img_pos_embed = nn.Parameter(torch.zeros(1, num_img_patches + 1, embed_dim))464 else:465 # logger.info("Use joint position embedding, for image and video we use same pos_embed.")466 self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))467 self.clip_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))468 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]469 # choose which layer to use checkpoint470 with_cp_list = [False] * depth471 if use_checkpoint:472 for idx in range(depth):473 if idx < checkpoint_num:474 with_cp_list[idx] = True475 # logger.info(f"Droppath rate: {dpr}")476 # logger.info(f"Checkpoint list: {with_cp_list}")477 478 self.blocks = nn.ModuleList([479 Block(embed_dim, num_heads, mlp_ratio, qkv_bias=qkv_bias,480 norm_layer=norm_layer_for_blocks,481 drop_path=dpr[i], init_values=init_values, attn_drop=0.,482 use_flash_attn=use_flash_attn, use_fused_mlp=use_fused_mlp,483 fused_mlp_heuristic=fused_mlp_heuristic,484 with_cp=with_cp_list[i],485 qk_normalization=qk_normalization,486 layerscale_no_force_fp32=layerscale_no_force_fp32,487 use_fused_rmsnorm=use_fused_rmsnorm)488 for i in range(depth)])489 self.clip_projector = AttentionPoolingBlock(490 dim=embed_dim, num_heads=attn_pool_num_heads, qkv_bias=True, qk_scale=None,491 drop=0., attn_drop=0., norm_layer=partial(nn.LayerNorm, eps=1e-5), out_dim=clip_embed_dim)492 493 # CLIP decoder494 self.clip_decoder = nn.ModuleList([495 Linear_Decoder(496 in_channels=embed_dim, 497 out_channels=clip_teacher_embed_dim, 498 norm_layer=partial(nn.LayerNorm, eps=1e-5), 499 clip_norm_type=clip_norm_type500 ) for _ in range(clip_return_layer)501 ])502 self.final_clip_decoder = nn.Identity()503 if clip_teacher_final_dim > 0:504 self.final_clip_decoder = Linear_Decoder(505 in_channels=clip_embed_dim, 506 out_channels=clip_teacher_final_dim, 507 norm_layer=partial(nn.LayerNorm, eps=1e-5), 508 clip_norm_type=clip_norm_type509 )510 511 self.init_pos_embed()512 trunc_normal_(self.cls_token, std=.02)513 self.apply(self._init_weights)514 self.fix_init_weight()515 516 def init_pos_embed(self):517 # logger.info("Init pos_embed from sincos pos_embed")518 if self.sep_pos_embed:519 raise NotImplementedError520 else:521 # trunc_normal_(self.pos_embed, std=.02)522 # trunc_normal_(self.clip_pos_embed, std=.02)523 pos_embed = get_3d_sincos_pos_embed(524 self.pos_embed.shape[-1], 525 self.patch_embed.grid_size[1], # height & weight526 self.patch_embed.grid_size[0], # t_size527 cls_token=True528 )529 self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))530 self.clip_pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))531 532 if self.sep_image_video_pos_embed:533 img_pos_embed = get_3d_sincos_pos_embed(534 self.pos_embed.shape[-1], 535 self.patch_embed.grid_size[1], # height & weight536 1,537 cls_token=True538 )539 self.img_pos_embed.data.copy_(torch.from_numpy(img_pos_embed).float().unsqueeze(0))540 self.clip_img_pos_embed.data.copy_(torch.from_numpy(img_pos_embed).float().unsqueeze(0))541 542 def _init_weights(self, m):543 if isinstance(m, nn.Linear):544 trunc_normal_(m.weight, std=.02)545 if isinstance(m, nn.Linear) and m.bias is not None:546 nn.init.constant_(m.bias, 0)547 elif isinstance(m, nn.LayerNorm):548 nn.init.constant_(m.bias, 0)549 nn.init.constant_(m.weight, 1.0)550 551 def fix_init_weight(self):552 def rescale(param, layer_id):553 param.div_(math.sqrt(2.0 * layer_id))554 555 for layer_id, layer in enumerate(self.blocks):556 rescale(layer.attn.proj.weight.data, layer_id + 1)557 rescale(layer.mlp.fc2.weight.data, layer_id + 1)558 559 @property560 def dtype(self):561 return self.patch_embed.proj.weight.dtype562 563 def get_num_layers(self):564 return len(self.blocks)565 566 @torch.jit.ignore567 def no_weight_decay(self):568 return {569 'pos_embed', 570 'pos_embed_spatial', 571 'pos_embed_temporal', 572 'pos_embed_cls',573 'img_pos_embed',574 'cls_token',575 'clip_pos_embed', 576 'clip_pos_embed_spatial', 577 'clip_pos_embed_temporal', 578 'clip_pos_embed_cls',579 'clip_img_pos_embed'580 }581 582 # @torch.cuda.amp.autocast(enabled=False)583 def forward(self, x, mask=None, use_image=False, x_vis_return_idx=-1, x_vis_only=False):584 # print(0, x.shape)585 x = self.patch_embed(x.type(self.dtype))586 # print(f"x.shape: {x.shape} x.dtype: {x.dtype}, model.dtype: {self.dtype}")587 B, T, L, C = x.shape # T: temporal; L: spatial588 x = x.view([B, T * L, C]) # (B, T * L, C)589 590 # append cls token591 cls_tokens = self.cls_token.expand(B, -1, -1)592 x = torch.cat((cls_tokens, x), dim=1) # (B, T * L + 1, C)593 # print(1, x.shape)594 595 # add pos_embed596 if self.sep_pos_embed:597 raise NotImplementedError598 else:599 if use_image:600 # print('use image') # No.601 if self.sep_image_video_pos_embed:602 pos_embed = self.img_pos_embed603 else:604 # (1, num_img_patches + 1, embed_dim)605 # print('origin pos_embed.shape:', self.pos_embed.shape)606 cls_pos_embed = self.pos_embed[:, 0:1, :]607 # print('cls_pos_embed.shape:', cls_pos_embed.shape)608 609 img_pos_embed = self.pos_embed[:, 1:, :].view(1, self.num_frames, self.patch_embed.num_patches // self.num_frames, self.embed_dim).mean(dim=1)610 # print('img_pos_embed.shape:', img_pos_embed.shape)611 612 pos_embed = torch.cat([cls_pos_embed, img_pos_embed], dim=1)613 # print('final img_pos_embed.shape:', pos_embed.shape)614 else:615 pos_embed = self.pos_embed616 pos_embed = pos_embed[:, :x.shape[1], :]617 x = x + pos_embed618 619 # mask tokens, ~mask means visible620 if mask is not None:621 x = x[~mask].reshape(B, -1, C) 622 else:623 x = x.reshape(B, -1, C) 624 residual = None625 x_clip = []626 for idx, blk in enumerate(self.blocks):627 if isinstance(x, tuple) and len(x) == 2:628 x, residual = x629 # print(f"\033[31m这是{idx}, {x.shape}\033[0m")630 x = blk(x, residual=residual)631 # return intermediate features632 if idx in self.return_index:633 if isinstance(x, tuple) and len(x) == 2:634 tmp_x, tmp_residual = x635 if residual is not None:636 x_clip.append(tmp_x + tmp_residual)637 else:638 x_clip.append(x)639 if idx == (self.depth + x_vis_return_idx):640 # print(f'idx = {idx} len(self.blocks)={len(self.blocks)}')641 break642 643 if isinstance(x, tuple) and len(x) == 2:644 x, residual = x645 if residual is not None:646 x = x + residual647 648 x_vis = x649 # print(f'x_vis.shape:{x_vis.shape}')650 if x_vis_only:651 return x_vis652 653 x_pool_vis = self.clip_projector(x_vis) 654 x_align = self.final_clip_decoder(x_pool_vis)655 # print(3, x_pool_vis.shape)656 # print(4, x_align.shape)657 658 # align CLIP659 x_clip = torch.stack(x_clip)660 K, B, _, C_CLIP = x_clip.shape661 # print(5, x_clip.shape)662 # add pos_embed663 if self.sep_pos_embed: 664 raise NotImplementedError665 else:666 if use_image:667 if self.sep_image_video_pos_embed:668 clip_pos_embed = self.clip_img_pos_embed669 else:670 # (1, num_img_patches + 1, embed_dim)671 # print('origin pos_embed.shape:', self.pos_embed.shape)672 clip_cls_pos_embed = self.clip_pos_embed[:, 0:1, :]673 # print('cls_pos_embed.shape:', cls_pos_embed.shape)674 675 clip_img_pos_embed = self.clip_pos_embed[:, 1:, :].view(1, self.num_frames, self.patch_embed.num_patches // self.num_frames, self.embed_dim).mean(dim=1)676 # print('img_pos_embed.shape:', img_pos_embed.shape)677 678 clip_pos_embed = torch.cat([clip_cls_pos_embed, clip_img_pos_embed], dim=1)679 # print('final img_pos_embed.shape:', pos_embed.shape)680 681 else:682 clip_pos_embed = self.clip_pos_embed683 684 clip_pos_embed = clip_pos_embed.repeat(B, 1, 1)685 if mask is not None:686 x_clip = x_clip + clip_pos_embed[~mask].view(B, -1, C_CLIP).unsqueeze(0).repeat(K, 1, 1, 1)687 else:688 clip_pos_embed = clip_pos_embed.unsqueeze(0).repeat(K, 1, 1, 1)689 clip_pos_embed = clip_pos_embed[:, :, :x_clip.shape[2], :]690 x_clip = x_clip + clip_pos_embed691 692 # CLIP decoder693 x_clip_align = []694 for idx, clip_decoder in enumerate(self.clip_decoder):695 x_clip_align.append(clip_decoder(x_clip[idx]))696 x_clip_align = torch.stack(x_clip_align)697 698 # print(f'x_vis.shape:{x_vis.shape}, x_pool_vis.shape:{x_pool_vis.shape}')699 return x_vis, x_pool_vis, x_clip_align, x_align700 701 702def pretrain_internvideo2_1b_patch14_224(config):703 # print(config.vision_encoder.num_frames)704 model = PretrainInternVideo2(705 in_chans=3, img_size=224, patch_size=14,706 embed_dim=1408, depth=40, num_heads=16, mlp_ratio=48/11,707 clip_embed_dim=config.vision_encoder.clip_embed_dim,708 attn_pool_num_heads=16, qkv_bias=False,709 drop_path_rate=0.25,710 init_values=0.00001,711 qk_normalization=True,712 use_flash_attn=config.vision_encoder.get('use_flash_attn', True),713 use_fused_rmsnorm=config.vision_encoder.get('use_fused_rmsnorm', True),714 use_fused_mlp=config.vision_encoder.get('use_fused_mlp', True),715 fused_mlp_heuristic=1,716 layerscale_no_force_fp32=False,717 num_frames=config.vision_encoder.num_frames,718 tubelet_size=config.vision_encoder.tubelet_size,719 sep_pos_embed=False,720 sep_image_video_pos_embed=config.vision_encoder.sep_image_video_pos_embed,721 use_checkpoint=config.vision_encoder.use_checkpoint,722 checkpoint_num=config.vision_encoder.checkpoint_num,723 clip_teacher_embed_dim=config.vision_encoder.clip_teacher_embed_dim,724 clip_teacher_final_dim=config.vision_encoder.clip_teacher_final_dim,725 clip_norm_type=config.vision_encoder.clip_norm_type,726 clip_return_layer=config.vision_encoder.clip_return_layer,727 clip_student_return_interval=config.vision_encoder.clip_student_return_interval,728 )729 730 if config.vision_encoder.pretrained is not None:731 # logger.info(f"Loading pretrained weights from {config.vision_encoder.pretrained}")732 state_dict = torch.load(config.vision_encoder.pretrained, map_location='cpu')733 interpolate_pos_embed_internvideo2(state_dict, model, orig_t_size=8)734 message = model.load_state_dict(state_dict, strict=False)735 # logger.info(message)736 else:737 pass738 # logger.info("No pretrained weights!!!")739 return model740 741 742 743def pretrain_internvideo2_6b_patch14_224(config):744 model = PretrainInternVideo2(745 in_chans=3, img_size=224, patch_size=14,746 embed_dim=3200, depth=48, num_heads=25, mlp_ratio=4,747 clip_embed_dim=config.vision_encoder.clip_embed_dim,748 attn_pool_num_heads=16, qkv_bias=False,749 drop_path_rate=0.3,750 init_values=0.00001,751 qk_normalization=True,752 use_flash_attn=config.vision_encoder.get('use_flash_attn', True),753 use_fused_rmsnorm=config.vision_encoder.get('use_fused_rmsnorm', True),754 use_fused_mlp=config.vision_encoder.get('use_fused_mlp', True),755 fused_mlp_heuristic=1,756 layerscale_no_force_fp32=False,757 num_frames=config.vision_encoder.num_frames,758 tubelet_size=config.vision_encoder.tubelet_size,759 sep_pos_embed=False,760 sep_image_video_pos_embed=config.vision_encoder.sep_image_video_pos_embed,761 use_checkpoint=config.vision_encoder.use_checkpoint,762 checkpoint_num=config.vision_encoder.checkpoint_num,763 clip_teacher_embed_dim=config.vision_encoder.clip_teacher_embed_dim,764 clip_teacher_final_dim=config.vision_encoder.clip_teacher_final_dim,765 clip_norm_type=config.vision_encoder.clip_norm_type,766 clip_return_layer=config.vision_encoder.clip_return_layer,767 clip_student_return_interval=config.vision_encoder.clip_student_return_interval,768 )769 770 if config.vision_encoder.pretrained is not None:771 # logger.info(f"Loading pretrained weights from {config.vision_encoder.pretrained}")772 state_dict = torch.load(config.vision_encoder.pretrained, map_location='cpu')773 interpolate_pos_embed_internvideo2(state_dict, model, orig_t_size=8)774 msg = model.load_state_dict(state_dict, strict=False)775 # logger.info(msg)776 else:777 pass778 # logger.info("No pretrained weights!!!")779 return model780 781