From cfac34ca6a23393dd19815e0f1454e3ecc48a6c8 Mon Sep 17 00:00:00 2001 From: forXuyx Date: Mon, 17 Mar 2025 19:49:29 +0800 Subject: [PATCH] =?UTF-8?q?=E4=B8=8Evlm=E7=BB=9F=E4=B8=80=E6=B5=81?= =?UTF-8?q?=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- model/dataset.py | 41 +++++++++++++++++---------- model/model_t2i.py | 66 ++++++++++++++++++++++++++++--------------- train_pretrain_t2i.py | 11 ++++---- 3 files changed, 76 insertions(+), 42 deletions(-) diff --git a/model/dataset.py b/model/dataset.py index 1984e23..e07eab9 100644 --- a/model/dataset.py +++ b/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('assistant\n', add_special_tokens=False).input_ids self.eos_id = tokenizer('\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 = [] - for image_name in image_paths.split(','): - image_name = image_name.strip() - image = Image.open(f'{self.images_path}/{image_name}') - image_tensor = self.transform(image) - image_tensors.append(image_tensor) - image_tensors = torch.stack(image_tensors, dim=0) + if self.img_pre_process: + image_emb_path = image_paths.replace('.jpg', '_emb.npy') + image_emb = np.load(f'{self.images_path}/{image_emb_path}') + image_emb = torch.tensor(image_emb, dtype=torch.float32) + # 加载预处理的图像token + image_token_path = image_paths.replace('.jpg', '_token.npy') + 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 diff --git a/model/model_t2i.py b/model/model_t2i.py index 6038f41..4144406 100644 --- a/model/model_t2i.py +++ b/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: - pixel_tensors = pixel_tensors.squeeze(2) - bs, num, c, im_h, im_w = pixel_tensors.shape - stack_dim = 1 if bs > 1 else 0 - # 获取图片的embedding - vision_tensors = torch.stack([ - MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0] - 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]) - # 获取图片的token - vision_tokens = torch.stack([ - MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1] - for i in range(num) - ], dim=stack_dim) - # 替换target中的'@......'为图片的token - target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1]) + if isinstance(pixel_tensors, torch.Tensor): + pixel_tensors = pixel_tensors.to(h.device) + if len(pixel_tensors.shape) == 6: + pixel_tensors = pixel_tensors.squeeze(2) + bs, num, c, im_h, im_w = pixel_tensors.shape + stack_dim = 1 if bs > 1 else 0 + # 获取图片的embedding + vision_tensors = torch.stack([ + MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0] + 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]) + # 获取图片的token + vision_tokens = torch.stack([ + MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1] + for i in range(num) + ], 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 = [] diff --git a/train_pretrain_t2i.py b/train_pretrain_t2i.py index b2ad0a0..70b6190 100644 --- a/train_pretrain_t2i.py +++ b/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, - image_special_token=model_config.image_special_token, - max_length=max_seq_len) + train_ds = T2IDataset(args.data_path, args.images_path, tokenizer, max_seq_len, args.img_process, + image_special_token=model_config.image_special_token) train_sampler = DistributedSampler(train_ds) if ddp else None train_loader = DataLoader( train_ds,