From 4a80ca7c7832c002f77f2e304f4ebd824aad5acc Mon Sep 17 00:00:00 2001
From: forXuyx <xuyuxiang22@mails.ucas.ac.cn>
Date: Sat, 22 Mar 2025 23:17:54 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E8=A7=86=E9=A2=91=E7=90=86?=
 =?UTF-8?q?=E8=A7=A3=E5=88=9D=E6=AD=A5=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 model/dataset.py          | 115 ++++++++++++++++++++
 model/model_videolm.py    | 132 +++++++++++++++++++++++
 train_pretrain_videolm.py | 214 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 461 insertions(+)
 create mode 100644 model/model_videolm.py
 create mode 100644 train_pretrain_videolm.py

diff --git a/model/dataset.py b/model/dataset.py
index dae04a7..84ccaf9 100644
--- a/model/dataset.py
+++ b/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
diff --git a/model/model_videolm.py b/model/model_videolm.py
new file mode 100644
index 0000000..edca4e3
--- /dev/null
+++ b/model/model_videolm.py
@@ -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
diff --git a/train_pretrain_videolm.py b/train_pretrain_videolm.py
new file mode 100644
index 0000000..eb90b0b
--- /dev/null
+++ b/train_pretrain_videolm.py
@@ -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)