与vlm统一流程

model/dataset.py

@@ -103,21 +103,18 @@ class VLMDataset(Dataset):
         return X, Y, loss_mask, image_tensors
 
 class T2IDataset(Dataset):
-    def __init__(self, jsonl_path, images_path, tokenizer, max_length=512,
+    def __init__(self, jsonl_path, images_path, tokenizer, max_length=512, img_pre_process=False,
                  image_special_token='@' * 256):
 
         super().__init__()
         self.samples = self.load_data(jsonl_path)
         self.images_path = images_path
 
+        self.img_pre_process = img_pre_process # 是否使用提前处理的图片
+
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.image_size = 256
-        self.transform = transforms.Compose([
-            transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, self.image_size)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-            ])
         self.image_token = image_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
@@ -172,14 +169,28 @@ class T2IDataset(Dataset):
 
         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)
+        loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) 
 
-        image_tensors = []
+        if self.img_pre_process:
-        for image_name in image_paths.split(','):
+            image_emb_path = image_paths.replace('.jpg', '_emb.npy')
-            image_name = image_name.strip()
+            image_emb = np.load(f'{self.images_path}/{image_emb_path}')
-            image = Image.open(f'{self.images_path}/{image_name}')
+            image_emb = torch.tensor(image_emb, dtype=torch.float32)
-            image_tensor = self.transform(image)
+            # 加载预处理的图像token
-            image_tensors.append(image_tensor)
+            image_token_path = image_paths.replace('.jpg', '_token.npy')
-        image_tensors = torch.stack(image_tensors, dim=0)   
+            image_token = np.load(f'{self.images_path}/{image_token_path}')
+            image_token = torch.tensor(image_token, dtype=torch.long)
+            return X, Y, loss_mask, (image_emb, image_token)
+        else:
+            image_tensors = []
+            for image_name in image_paths.split(','):
+                image_name = image_name.strip()
+                image = Image.open(f'{self.images_path}/{image_name}').convert("RGB")
+                image = center_crop_arr(image, self.image_size)
+                image = np.array(image) / 255.
+                image = 2.0 * image - 1.0
+                image = torch.tensor(image, dtype=torch.float32)
+                image_tensor = torch.einsum('hwc->chw', image)
+                image_tensors.append(image_tensor)
+            image_tensors = torch.stack(image_tensors, dim=0)   
 
-        return X, Y, loss_mask, image_tensors
+            return X, Y, loss_mask, image_tensors

model/model_t2i.py

@@ -3,13 +3,29 @@ from .model import *
 from typing import Optional, Tuple, List
 from torch import nn
 import warnings
-from model.model_vlm import MiniMindVLM, VisionProj
+from model.model_vlm import MiniMindVLM
 from model.model_vq import VQ_models
 import torch
 from einops import rearrange
 
 warnings.filterwarnings('ignore')
 
+class VisionProj(nn.Module):
+    def __init__(self, ve_dim=768, lm_dim=512, hidden_dim=256):
+        super().__init__()
+        self.ve_dim = ve_dim
+        self.lm_dim = lm_dim
+        self.hidden_dim = hidden_dim
+        self.vision_proj = nn.Sequential(
+            nn.Linear(self.ve_dim, self.hidden_dim),
+            nn.ReLU(),
+            nn.Linear(self.hidden_dim, self.lm_dim)
+        )
+
+    def forward(self, image_encoders):
+        vision_proj = self.vision_proj(image_encoders)
+        return vision_proj
+
 
 # 继承自语言模型
 class MiniMindT2I(MiniMindVLM):
@@ -20,14 +36,12 @@ class MiniMindT2I(MiniMindVLM):
         if not params: params = VLMConfig()
         self.params = params
         self.image_tokenizer = self.__class__.get_image_tokenizer()
-        # 这里重写是为了将output映射到codebook_size
-        self.output = nn.Linear(params.dim, 16384, bias=False)
         self.vision_proj = VisionProj(ve_dim=8, lm_dim=params.dim)
 
     @staticmethod
-    def get_image_tokenizer(model_path="./model/minimind_img_tokenizer/vq_ds16_t2i.pt"):
+    def get_image_tokenizer(model_path="./model/minimind_img_tokenizer/minimind_img_tokenizer.pt"):
         model = VQ_models['VQ-16'](
-            codebook_size=16384,
+            codebook_size=6400,
             codebook_embed_dim=8)
         
         # 加载模型
@@ -105,23 +119,31 @@ class MiniMindT2I(MiniMindVLM):
         h = self.tok_embeddings(input_ids)
 
         if pixel_tensors is not None and start_pos == 0:
-            if len(pixel_tensors.shape) == 6:
+            if isinstance(pixel_tensors, torch.Tensor):
-                pixel_tensors = pixel_tensors.squeeze(2)
+                pixel_tensors = pixel_tensors.to(h.device)
-            bs, num, c, im_h, im_w = pixel_tensors.shape
+                if len(pixel_tensors.shape) == 6:
-            stack_dim = 1 if bs > 1 else 0
+                    pixel_tensors = pixel_tensors.squeeze(2)
-            # 获取图片的embedding
+                bs, num, c, im_h, im_w = pixel_tensors.shape
-            vision_tensors = torch.stack([
+                stack_dim = 1 if bs > 1 else 0
-                MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
+                # 获取图片的embedding
-                for i in range(num)
+                vision_tensors = torch.stack([
-            ], dim=stack_dim)
+                    MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
-            h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
+                    for i in range(num)
-            # 获取图片的token
+                ], dim=stack_dim)
-            vision_tokens = torch.stack([
+                h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
-                MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
+                # 获取图片的token
-                for i in range(num)
+                vision_tokens = torch.stack([
-            ], dim=stack_dim)
+                    MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
-            # 替换target中的'@......'为图片的token
+                    for i in range(num)
-            target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
+                ], dim=stack_dim)
+                # 替换target中的'@......'为图片的token
+                target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
+            else:
+                vision_tensors, vision_tokens = pixel_tensors
+                vision_tensors = vision_tensors.to(h.device)
+                vision_tokens = vision_tokens.to(h.device)
+                h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
+                target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
 
         pos_cis = self.pos_cis[start_pos:start_pos + input_ids.shape[1]]
         past_kvs = []

train_pretrain_t2i.py

@@ -39,13 +39,14 @@ def train_epoch(epoch, wandb):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
-        pixel_tensors = pixel_tensors.to(args.device)
+        pixel_tensors = pixel_tensors
         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, Y, pixel_tensors=pixel_tensors)
+            Y = res.target_ids
             loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 res.target_ids.view(-1)
@@ -147,7 +148,7 @@ if __name__ == "__main__":
     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_t2i_data.jsonl")
-    parser.add_argument("--images_path", type=str, default="./dataset/pretrain_images")
+    parser.add_argument("--images_path", type=str, default="./dataset/pretrain_t2i_code")
     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)
@@ -159,6 +160,7 @@ if __name__ == "__main__":
     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)
+    parser.add_argument('--img_process', type=bool, default=True)
     args = parser.parse_args()
 
     model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len,
@@ -189,9 +191,8 @@ if __name__ == "__main__":
 
     model, tokenizer = init_model(model_config)
 
-    train_ds = T2IDataset(args.data_path, args.images_path, tokenizer,
+    train_ds = T2IDataset(args.data_path, args.images_path, tokenizer, max_seq_len, args.img_process,
-                          image_special_token=model_config.image_special_token,
+                          image_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,