Views
No views yet
1import os
2import torch
3import json
4import argparse
5from tqdm import tqdm
6from collections import defaultdict
7import torch.nn.functional as F
8from time import time
9from easydict import EasyDict as edict
10
11from model.mico import *
12
13
14def load_from_pretrained_dir(pretrain_dir, video_resolution=224, return_modal="full"):
15
16 checkpoint_dir = os.path.join(pretrain_dir,'ckpt')
17 file_cfg = edict(json.load(open(os.path.join(pretrain_dir,'log','hps.json'))))
18 model_cfg = file_cfg.model_cfg
19 checkpoint_ls = [ i for i in os.listdir(checkpoint_dir) if i.startswith('model_step')]
20 checkpoint_ls = [int(i.split('_')[2].split('.')[0]) for i in checkpoint_ls]
21 checkpoint_ls.sort()
22 step = checkpoint_ls[-1]
23
24 checkpoint_name = 'model_step_'+str(step)+'.pt'
25 ckpt_file = os.path.join(checkpoint_dir, checkpoint_name)
26 checkpoint = torch.load(ckpt_file, map_location = 'cpu')
27 print(f'load_from_pretrained: {ckpt_file}')
28
29 new_ckpt = {}
30 for k,v in checkpoint.items():
31 if 'video' in k:
32 new_ckpt[k.replace('video','vision')]=v
33 elif 'evaclip_model' in k:
34 new_ckpt[k.replace('evaclip_model','vision_encoder')]=v
35 elif 'clip_model' in k:
36 new_ckpt[k.replace('clip_model','vision_encoder')]=v
37 else:
38 new_ckpt[k] = v.float()
39
40 checkpoint = new_ckpt
41
42 if model_cfg.frame_embedding_type == 'adaptive':
43
44 if 'vision_frame_embedding' in checkpoint:
45 pretrain_embed = checkpoint['vision_frame_embedding']
46 if pretrain_embed.shape[1]!=model_cfg.max_vision_sample_num:
47 pretrain_embed = F.interpolate(pretrain_embed.permute(0,2,1),model_cfg.max_vision_sample_num,mode='nearest').permute(0,2,1)
48 checkpoint['vision_frame_embedding'] = pretrain_embed
49 else:
50 pretrain_embed = checkpoint['vision_perceiver.vision_frame_embedding']
51 if pretrain_embed.shape[1]!=model_cfg.max_vision_sample_num:
52 pretrain_embed = F.interpolate(pretrain_embed.permute(0,2,1),model_cfg.max_vision_sample_num,mode='nearest').permute(0,2,1)
53 checkpoint['vision_perceiver.vision_frame_embedding'] = pretrain_embed
54
55 if 'audio_frame_embedding' in checkpoint:
56 pretrain_embed_a = checkpoint['audio_frame_embedding']
57 if pretrain_embed_a.shape[1]!=model_cfg.max_audio_sample_num:
58 pretrain_embed_a = F.interpolate(pretrain_embed_a.permute(0,2,1),model_cfg.max_audio_sample_num,mode='nearest').permute(0,2,1)
59 checkpoint['audio_frame_embedding'] = pretrain_embed_a
60
61 if model_cfg.vision_encoder_type.startswith('clip'):
62 vision_width = checkpoint["vision_encoder.visual.positional_embedding"].shape[1]
63 vision_layers = len([k for k in checkpoint.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
64 vision_patch_size = checkpoint["vision_encoder.visual.conv1.weight"].shape[-1]
65
66 grid_size = round((checkpoint["vision_encoder.visual.positional_embedding"].shape[0] - 1) ** 0.5)
67
68 src = checkpoint["vision_encoder.visual.positional_embedding"]
69 src_cls = src[0:1]
70 src_oth = src[1:]
71 new_grid_size = model_cfg.vision_resolution // vision_patch_size
72 if new_grid_size!=grid_size:
73 src_oth = F.interpolate(src_oth.reshape(grid_size,grid_size,vision_width).permute(2,0,1).unsqueeze(0),(new_grid_size,new_grid_size),mode='bilinear')
74 src_oth = src_oth[0].permute(1,2,0).reshape(-1,src.shape[-1])
75 tgt = torch.cat((src_cls,src_oth),dim=0)
76 checkpoint["vision_encoder.visual.positional_embedding"] = tgt
77
78 elif model_cfg.vision_encoder_type.startswith('evaclip'):
79
80 vision_width = checkpoint["vision_encoder.visual.pos_embed"].shape[2]
81 vision_layers = len([k for k in checkpoint.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
82
83 vision_patch_size = checkpoint["vision_encoder.visual.patch_embed.proj.weight"].shape[-1]
84
85 grid_size = round((checkpoint["vision_encoder.visual.pos_embed"].shape[1] - 1) ** 0.5)
86
87 src = checkpoint["vision_encoder.visual.pos_embed"][0]
88 src_cls = src[0:1]
89 src_oth = src[1:]
90 new_grid_size = model_cfg.vision_resolution // vision_patch_size
91 if new_grid_size!=grid_size:
92 src_oth = F.interpolate(src_oth.reshape(grid_size,grid_size,vision_width).permute(2,0,1).unsqueeze(0),(new_grid_size,new_grid_size),mode='bilinear')
93 src_oth = src_oth[0].permute(1,2,0).reshape(-1,src.shape[-1])
94 tgt = torch.cat((src_cls,src_oth),dim=0)
95 checkpoint["vision_encoder.visual.pos_embed"] = tgt.unsqueeze(0)
96 else:
97 pass
98
99 if return_modal=="full":
100 new_ckpt = checkpoint
101 elif return_modal=="uni":
102 new_ckpt = defaultdict()
103 for k in checkpoint.keys():
104 if "video_encoder" in k:
105 new_k = ".".join(k.split(".")[1:])
106 new_ckpt[new_k] = checkpoint[k]
107 elif return_modal=="text":
108 new_ckpt = defaultdict()
109 for k in checkpoint.keys():
110 if "multimodal_encoder" in k:
111 new_k = ".".join(k.split(".")[1:])
112 new_ckpt[new_k] = checkpoint[k]
113 else:
114 pass
115
116 return new_ckpt, model_cfg
117
118
119if __name__ == "__main__":
120 # import ipdb
121 # ipdb.set_trace()
122 device = "cuda"
123 from model.imageprocessor import ImageProcessor
124 pretrain_path = 'MiCo-g' # please check your
125 checkpoint, opts = load_from_pretrained_dir("MiCo-g", video_resolution=224, return_modal="full")
126 model = MiCo.from_pretrained(opts,checkpoint).to(device)
127 image_file = "example/test.jpeg"
128 proc = ImageProcessor(image_resolution=224, image_encoder_type="swin", training=True)
129 image_input = proc(image_file).to(device)
130 image_input = image_input.unsqueeze(1) # image as a 1 frame video
131
132 video_output = model.forward_vision_encoder(image_input)
133 video_output_pooled = model.pool_vision_for_contra(video_output)
134 feat_v = model.contra_head_v(video_output_pooled)
135 feat_v = F.normalize(feat_v,dim=-1)
136
137 texts = ["a man is skiing in a snowy day.", "it's a hot day"]
138 caption_tokens = model.multimodal_encoder.tokenizer(texts,
139 padding="max_length",
140 truncation=True,
141 max_length=30,
142 return_tensors="pt")
143 caption_tokens = caption_tokens.to(torch.device('cuda'))
144 input_ids = caption_tokens.input_ids
145 attention_mask = caption_tokens.attention_mask
146 caption_output = model.forward_multimodal_encoder(input_ids, attention_mask).sequence_output
147 caption_output_pooled = model.pool_text_for_contra(caption_output)
148 feat_t = model.contra_head_t(caption_output_pooled)
149 feat_t = F.normalize(feat_t,dim=-1)
150
151
152 sim_t2v = torch.matmul(feat_t, feat_v.permute(1,0))
153 print(sim_t2v)
154
155 video_input = model.get_multimodal_forward_input_vision(video_output)
156 slice_output = model.forward_multimodal_encoder(input_ids, attention_mask, video_input).sequence_output
157 slice_scores = F.softmax(model.itm_head(slice_output[:,0]),dim=1)[:,1]
158 print(slice_scores)
159
160
161 video_input = model.get_multimodal_forward_input_vision(video_output)
162 init_input_ids = torch.ones(video_input.size(0), 1).long().cuda().fill_(model.multimodal_encoder.tokenizer.bos_token_id)
163 init_attention_mask = init_input_ids.new_ones(video_input.size(0), 1, 1)
164 outputs = model.multimodal_encoder.generate(input_ids=init_input_ids,
165 attention_mask=init_attention_mask,
166 encoder_hidden_states=video_input,
167 max_new_tokens=model.max_caption_len,
168 num_beams=model.beam_size,
169 eos_token_id=model.multimodal_encoder.tokenizer.sep_token_id,
170 pad_token_id=model.multimodal_encoder.tokenizer.pad_token_id,
171 length_penalty=0.6)
172 outputs_newgen = outputs[:,1:]
173 captions = model.multimodal_encoder.tokenizer.batch_decode(outputs_newgen, skip_special_tokens=True)
174 print(captions)
knowledge modality and interface modality. Knowledge modalities, primarily derived from raw sensors, contribute knowledge in diverse formats. For example, images and depth maps offer visual knowledge, while audio and video provide auditory and spatiotemporal knowledge. The language modality, developed by humans, is inherently more abstract and naturally functions as the interface modality, facilitating learning, reasoning, and the coordination of knowledge. To this end, we design an omni-modal learning architecture, illustrated in Figure (b), with two distinct branches: one for knowledge modalities and one for the interface modality, i.e. natural language. The knowledge and interface modalities are aligned through a novel generative reasoning method.
1pip install gdown
2gdown 1AIQjV1KU8K4OXiO-4gFirxkoxt3twWIq --folder
3python inference_demo.py@article{zhang2024explore,
title={Explore the Limits of Omni-modal Pretraining at Scale},
author={Zhang, Yiyuan and Li, Handong and Liu, Jing and Yue, Xiangyu},
journal={arXiv preprint arXiv:2406.xxxxx},
year={2024}
}