File size: 3,801 Bytes
edf2ce7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import logging
import json
import torch
from torch import nn
from .config import InternVideo2Config, EasyDict
from .internvideo2 import pretrain_internvideo2_1b_patch14_224, pretrain_internvideo2_6b_patch14_224
from transformers.utils import logging
import warnings
warnings.filterwarnings("ignore")
class InternVideo2_Stage2(nn.Module):
"""docstring for InternVideo2_Stage2"""
def __init__(self, config, is_pretrain=True):
super(InternVideo2_Stage2, self).__init__()
# if isinstance(config, InternVideo2Config):
# config_str = str(config)
# config_str = config_str.replace('InternVideo2Config ', '')
# config_json = json.loads(config_str)
# config = EasyDict(config_json)
# self.config = config
self.config = config
self.is_pretrain = is_pretrain
self.vision_width = config.model.vision_encoder.clip_embed_dim
# self.text_width = config.model.text_encoder.d_model
self.embed_dim = config.model.embed_dim
# create modules.
self.vision_encoder = self.build_vision_encoder()
if config.model.get("freeze_vision", False):
self.freeze_vision()
self.vision_proj = nn.Linear(self.vision_width, self.embed_dim)
self.temp = nn.parameter.Parameter(torch.ones([]) * config.model.temp)
self.uta_image_only = config.criterion.get('uta_image_only', False)
# logger.info(f"uta_image_only={self.uta_image_only}")
def freeze_vision(self):
"""freeze vision encoder"""
for p in self.vision_encoder.parameters():
p.requires_grad = False
def no_weight_decay(self):
ret = {"temp"}
ret.update(
{"vision_encoder." + k for k in self.vision_encoder.no_weight_decay()}
)
# ret.update(
# {"text_encoder." + k for k in self.text_encoder.no_weight_decay()}
# )
return ret
@property
def dtype(self):
return self.vision_encoder.patch_embed.proj.weight.dtype
def encode_vision(self, image):
"""encode image / videos as features.
Args:
image (torch.Tensor): The input images. Shape(B, N, C, H, W)
test (bool): Whether testing.
Returns: tuple.
- vision_embeds (torch.Tensor): The output features. Shape: [B,N,C].
- pooled_vision_embeds (torch.Tensor): The pooled output features. Shape: [B,1,C].
- student_output (torch.Tensor): The features of alignment. Shape: [K,B,N,C].
- clip_output (torch.Tensor): The features of clip. Shape: [K,B,N,C].
"""
T = image.shape[1]
use_image = True if T == 1 else False
image = image.permute(0, 2, 1, 3, 4) # [B,N,C,H,W] -> [B,C,N,H,W]
# whether save temporal dimension
# keep_temporal=self.config.model.vision_encoder.keep_temporal
vision_embeds, pooled_vision_embeds, _, _ = self.vision_encoder(
image, None, use_image)
return vision_embeds, pooled_vision_embeds
def build_vision_encoder(self):
"""build vision encoder
Returns: (vision_encoder, clip_teacher). Each is a `nn.Module`.
"""
encoder_name = self.config.model.vision_encoder.name
# logger.info(f"Build vision_encoder: {encoder_name}")
if encoder_name == 'pretrain_internvideo2_1b_patch14_224':
vision_encoder = pretrain_internvideo2_1b_patch14_224(self.config.model)
elif encoder_name == 'pretrain_internvideo2_6b_patch14_224':
vision_encoder = pretrain_internvideo2_6b_patch14_224(self.config.model)
else:
raise ValueError(f"Not implemented: {encoder_name}")
return vision_encoder
|