实现视频理解初步代码
forXuyx
parent
c1a069b532
commit
4a80ca7c78
115
model/dataset.py
115
model/dataset.py
|
|
@ -4,10 +4,48 @@ from torch.utils.data import Dataset, DataLoader
|
|||
import torch
|
||||
from .model_vlm import MiniMindVLM
|
||||
import os
|
||||
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
||||
|
||||
def video2image(video_path, num_frames=8, size=224):
|
||||
def preprocess(size, n_px):
|
||||
return Compose([
|
||||
Resize(size, interpolation=InterpolationMode.BICUBIC),
|
||||
CenterCrop(size),
|
||||
lambda image: image.convert("RGB"),
|
||||
ToTensor(),
|
||||
Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
|
||||
])(n_px)
|
||||
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
cap = cv2.VideoCapture(video_path, cv2.CAP_FFMPEG)
|
||||
frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
||||
|
||||
if fps < 1 or frameCount < 1:
|
||||
images = np.zeros([3, size, size], dtype=np.float32)
|
||||
print("ERROR: problem reading video file: ", video_path)
|
||||
else:
|
||||
frames_idx = np.sort(np.random.choice(frameCount, min(num_frames, frameCount), replace=False))
|
||||
|
||||
images = np.zeros([len(frames_idx), 3, size, size], dtype=np.float32)
|
||||
|
||||
for i, idx in enumerate(frames_idx):
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
continue
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
images[i,:,:,:] = preprocess(size, Image.fromarray(frame).convert("RGB"))
|
||||
|
||||
cap.release()
|
||||
video_frames = torch.tensor(images)
|
||||
return video_frames
|
||||
|
||||
class VLMDataset(Dataset):
|
||||
def __init__(self, jsonl_path, images_path, tokenizer, preprocess=None, max_length=512,
|
||||
image_special_token='@' * 196):
|
||||
|
|
@ -84,3 +122,80 @@ class VLMDataset(Dataset):
|
|||
image_tensors = torch.stack(image_tensors, dim=0)
|
||||
|
||||
return X, Y, loss_mask, image_tensors
|
||||
|
||||
|
||||
class VideoLMDataset(Dataset):
|
||||
def __init__(self, jsonl_path, videos_path, tokenizer, max_length=512,
|
||||
video_special_token='@' * 196):
|
||||
|
||||
super().__init__()
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.videos_path = videos_path
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.video_token = video_special_token
|
||||
self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def load_data(self, path):
|
||||
samples = []
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
for line_num, line in enumerate(f, 1):
|
||||
data = json.loads(line.strip())
|
||||
samples.append(data)
|
||||
return samples
|
||||
|
||||
def _create_chat_prompt(self, conversations):
|
||||
messages = []
|
||||
for i, turn in enumerate(conversations):
|
||||
role = 'user' if i % 2 == 0 else 'assistant'
|
||||
messages.append({"role": role, "content": turn['content'].replace('<video>', self.video_token)})
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=False
|
||||
)
|
||||
|
||||
def _generate_loss_mask(self, input_ids):
|
||||
loss_mask = [0] * len(input_ids)
|
||||
i = 0
|
||||
while i < len(input_ids):
|
||||
if input_ids[i:i + len(self.bos_id)] == self.bos_id:
|
||||
start = i + len(self.bos_id)
|
||||
end = start
|
||||
while end < len(input_ids):
|
||||
if input_ids[end:end + len(self.eos_id)] == self.eos_id:
|
||||
break
|
||||
end += 1
|
||||
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
|
||||
loss_mask[j] = 1
|
||||
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
|
||||
else:
|
||||
i += 1
|
||||
return loss_mask
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
sample = self.samples[index]
|
||||
video_paths = sample['video']
|
||||
prompt = self._create_chat_prompt(sample['conversations'])
|
||||
input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
|
||||
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
|
||||
loss_mask = self._generate_loss_mask(input_ids)
|
||||
|
||||
X = torch.tensor(input_ids[:-1], dtype=torch.long)
|
||||
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
||||
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
|
||||
|
||||
video_tensors = []
|
||||
for video_name in video_paths.split(','):
|
||||
video_name = video_name.strip()
|
||||
video_path = f'{self.videos_path}/{video_name}'
|
||||
video_tensor = video2image(video_path)
|
||||
video_tensors.append(video_tensor)
|
||||
video_tensors = torch.stack(video_tensors, dim=0)
|
||||
|
||||
return X, Y, loss_mask, video_tensors
|
||||
|
|
|
|||
|
|
@ -0,0 +1,132 @@
|
|||
from .VLMConfig import VLMConfig
|
||||
from .model import *
|
||||
from typing import Optional, Tuple, List
|
||||
from torch import nn
|
||||
import warnings
|
||||
from transformers import CLIPProcessor, CLIPModel
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
class VisionProj(nn.Module):
|
||||
def __init__(self, ve_dim=768, lm_dim=512):
|
||||
super().__init__()
|
||||
self.ve_dim = ve_dim
|
||||
self.lm_dim = lm_dim
|
||||
self.vision_proj = nn.Sequential(
|
||||
nn.Linear(self.ve_dim, self.lm_dim)
|
||||
)
|
||||
|
||||
def forward(self, image_encoders):
|
||||
vision_proj = self.vision_proj(image_encoders)
|
||||
return vision_proj
|
||||
|
||||
|
||||
# 继承自语言模型
|
||||
class MiniMindVideoLM(MiniMindLM):
|
||||
config_class = VLMConfig
|
||||
|
||||
def __init__(self, params: VLMConfig = None):
|
||||
super().__init__(params)
|
||||
if not params: params = VLMConfig()
|
||||
self.params = params
|
||||
self.vision_encoder = self.__class__.get_vision_model()
|
||||
self.vision_proj = VisionProj(lm_dim=params.dim)
|
||||
|
||||
@staticmethod
|
||||
def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"):
|
||||
model = CLIPModel.from_pretrained(model_path)
|
||||
# 冻结 vision_encoder 的所有参数
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
return model.eval()
|
||||
|
||||
@staticmethod
|
||||
def get_video_embeddings(video_tensors, vision_model):
|
||||
# video_tensors: (bs, frame, c, h, w)
|
||||
vid_embeddings = []
|
||||
for video_tensor in video_tensors:
|
||||
with torch.no_grad():
|
||||
outputs = vision_model.vision_model(pixel_values=video_tensor)
|
||||
vid_embedding = outputs.last_hidden_state[:, 1:, :].squeeze()
|
||||
vid_embedding = vid_embedding / vid_embedding.norm(dim=-1, keepdim=True)
|
||||
vid_embedding = torch.mean(vid_embedding, dim=0)
|
||||
vid_embedding = vid_embedding / vid_embedding.norm(dim=-1, keepdim=True)
|
||||
vid_embeddings.append(vid_embedding)
|
||||
vid_embeddings = torch.stack(vid_embeddings, dim=0)
|
||||
|
||||
return vid_embeddings
|
||||
|
||||
def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
|
||||
def find_indices(tokens, image_ids):
|
||||
image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
|
||||
len_image_ids = len(image_ids)
|
||||
if len_image_ids > tokens.size(1):
|
||||
return None
|
||||
tokens_view = tokens.unfold(1, len_image_ids, 1)
|
||||
matches = (tokens_view == image_ids_tensor).all(dim=2)
|
||||
return {
|
||||
batch_idx: [(idx.item(), idx.item() + len_image_ids - 1) for idx in
|
||||
matches[batch_idx].nonzero(as_tuple=True)[0]]
|
||||
for batch_idx in range(tokens.size(0)) if matches[batch_idx].any()
|
||||
} or None
|
||||
|
||||
image_indices = find_indices(tokens, self.params.image_ids)
|
||||
if vision_tensors is not None and image_indices:
|
||||
vision_proj = self.vision_proj(vision_tensors)
|
||||
if len(vision_proj.shape) == 3:
|
||||
vision_proj = vision_proj.unsqueeze(0)
|
||||
new_h = []
|
||||
for i in range(h.size(0)):
|
||||
if i in image_indices:
|
||||
h_i = h[i]
|
||||
img_idx = 0
|
||||
for start_idx, end_idx in image_indices[i]:
|
||||
if img_idx < vision_proj.size(1):
|
||||
h_i = torch.cat((h_i[:start_idx], vision_proj[i][img_idx], h_i[end_idx + 1:]), dim=0)[
|
||||
:seqlen]
|
||||
img_idx += 1
|
||||
new_h.append(h_i)
|
||||
else:
|
||||
new_h.append(h[i])
|
||||
return torch.stack(new_h, dim=0)
|
||||
return h
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
**args):
|
||||
start_pos = args.get('start_pos', 0)
|
||||
pixel_tensors = args.get('pixel_tensors', None)
|
||||
h = self.tok_embeddings(input_ids)
|
||||
|
||||
if pixel_tensors is not None and start_pos == 0:
|
||||
if len(pixel_tensors.shape) == 7:
|
||||
pixel_tensors = pixel_tensors.squeeze(2)
|
||||
bs, num, frame, c, im_h, im_w = pixel_tensors.shape
|
||||
stack_dim = 1 if bs > 1 else 0
|
||||
vision_tensors = torch.stack([
|
||||
MiniMindVideoLM.get_video_embeddings(pixel_tensors[:, i, :, :, :, :], self.vision_encoder)
|
||||
for i in range(num)
|
||||
], dim=stack_dim)
|
||||
h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
|
||||
|
||||
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.shape[1]]
|
||||
past_kvs = []
|
||||
for l, layer in enumerate(self.layers):
|
||||
h, past_kv = layer(
|
||||
h, pos_cis,
|
||||
past_key_value=past_key_values[l] if past_key_values else None,
|
||||
use_cache=use_cache
|
||||
)
|
||||
past_kvs.append(past_kv)
|
||||
|
||||
logits = self.output(self.norm(h))
|
||||
aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
|
||||
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
return self.OUT
|
||||
|
|
@ -0,0 +1,214 @@
|
|||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import json
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
from model.model_videolm import MiniMindVideoLM
|
||||
from model.VLMConfig import VLMConfig
|
||||
from model.dataset import VideoLMDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def Logger(content):
|
||||
if not ddp or dist.get_rank() == 0:
|
||||
print(content)
|
||||
|
||||
|
||||
def get_lr(current_step, total_steps, lr):
|
||||
return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
|
||||
|
||||
|
||||
def train_epoch(epoch, wandb):
|
||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||
start_time = time.time()
|
||||
for step, (X, Y, loss_mask, pixel_tensors) in enumerate(train_loader):
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
pixel_tensors = pixel_tensors.to(args.device)
|
||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
with ctx:
|
||||
res = model(X, pixel_tensors=pixel_tensors)
|
||||
loss = loss_fct(
|
||||
res.logits.view(-1, res.logits.size(-1)),
|
||||
Y.view(-1)
|
||||
).view(Y.size())
|
||||
|
||||
loss = (loss * loss_mask).sum() / loss_mask.sum()
|
||||
loss += res.aux_loss
|
||||
loss = loss / args.accumulation_steps
|
||||
|
||||
scaler.scale(loss).backward()
|
||||
|
||||
if (step + 1) % args.accumulation_steps == 0:
|
||||
scaler.unscale_(optimizer)
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
||||
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
|
||||
epoch + 1,
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if model_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.dim}{moe_path}.pth'
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
clean_state_dict = {
|
||||
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
|
||||
}
|
||||
torch.save(clean_state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(model_config: VLMConfig):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
moe_path = '_moe' if model_config.use_moe else ''
|
||||
# 加载纯语言模型权重
|
||||
ckp = f'./out/lm_{model_config.dim}{moe_path}.pth'
|
||||
model = MiniMindVideoLM(model_config)
|
||||
# state_dict = torch.load(ckp, map_location=args.device)
|
||||
# model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
# 冻结除 vision_proj 外的所有参数
|
||||
for name, param in model.named_parameters():
|
||||
if 'vision_proj' not in name:
|
||||
param.requires_grad = False
|
||||
# 可训练
|
||||
if hasattr(model, "layers"):
|
||||
last_two_layers = model.layers[-1:]
|
||||
for layer in last_two_layers:
|
||||
for param in layer.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
|
||||
return model.to(args.device), tokenizer
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if not ddp: return
|
||||
global ddp_local_rank, DEVICE
|
||||
|
||||
dist.init_process_group(backend="nccl")
|
||||
ddp_rank = int(os.environ["RANK"])
|
||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
||||
ddp_world_size = int(os.environ["WORLD_SIZE"])
|
||||
DEVICE = f"cuda:{ddp_local_rank}"
|
||||
torch.cuda.set_device(DEVICE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--epochs", type=int, default=4)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--learning_rate", type=float, default=4e-4)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", default=False, action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
|
||||
parser.add_argument("--num_workers", type=int, default=8)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/pretrain_video.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="./dataset/MSRVTT/videos/all")
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
parser.add_argument("--warmup_iters", type=int, default=0)
|
||||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=640, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len)
|
||||
max_seq_len = model_config.max_seq_len
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
tokens_per_iter = args.batch_size * max_seq_len
|
||||
torch.manual_seed(1337)
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-V Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
|
||||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
|
||||
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
|
||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||||
if ddp:
|
||||
init_distributed_mode()
|
||||
args.device = torch.device(DEVICE)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
wandb = None
|
||||
|
||||
model, tokenizer = init_model(model_config)
|
||||
|
||||
train_ds = VideoLMDataset(args.data_path, args.images_path, tokenizer,
|
||||
video_special_token=model_config.image_special_token,
|
||||
max_length=max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
batch_size=args.batch_size,
|
||||
pin_memory=True,
|
||||
drop_last=False,
|
||||
shuffle=False,
|
||||
num_workers=args.num_workers,
|
||||
sampler=train_sampler
|
||||
)
|
||||
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
|
||||
optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)
|
||||
|
||||
if ddp:
|
||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
for epoch in range(args.epochs):
|
||||
train_epoch(epoch, wandb)
|
||||
Loading…
Reference in New Issue