diff --git a/eval_t2i.py b/eval_t2i.py new file mode 100644 index 0000000..6db11a7 --- /dev/null +++ b/eval_t2i.py @@ -0,0 +1,93 @@ +import argparse +import os +import random +import numpy as np +import torch +import warnings +import torch.nn.functional as F +from PIL import Image +from transformers import AutoTokenizer, AutoModelForCausalLM +from model.model_t2i import MiniMindT2I +from model.VLMConfig import VLMConfig +from transformers import logging as hf_logging + +hf_logging.set_verbosity_error() + +warnings.filterwarnings('ignore') + + +def count_parameters(model): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def init_model(lm_config, device): + tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') + moe_path = '_moe' if args.use_moe else '' + ckp = f'./{args.out_dir}/sft_t2i_{args.dim}{moe_path}.pth' + model = MiniMindT2I(lm_config) + state_dict = torch.load(ckp, map_location=device) + model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False) + + print(f'T2I参数量:{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') + + return model.eval().to(device), tokenizer + + +def setup_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Chat with MiniMind") + parser.add_argument('--lora_name', default='None', type=str) + parser.add_argument('--out_dir', default='out', type=str) + parser.add_argument('--temperature', default=0.65, type=float) + parser.add_argument('--top_p', default=0.85, type=float) + parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str) + # MiniMind2-Small (26M):(dim=512, n_layers=8) + # MiniMind2 (104M):(dim=768, n_layers=16) + 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() + + lm_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe) + + model, tokenizer = init_model(lm_config, args.device) + + + def chat_with_vlm(prompt): + messages = [{"role": "user", "content": prompt}] + + new_prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True + )[-args.max_seq_len + 1:] + + with torch.no_grad(): + x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0) + outputs = model.generate( + x, + temperature=args.temperature, + top_p=args.top_p, + pad_token_id=tokenizer.pad_token_id, + ) + image = model.image_tokenizer.decode_code(outputs, (1, 8, 16, 16)) + image = F.interpolate(image, size=[256, 256], mode='bicubic').permute(0, 2, 3, 1)[0] + image = torch.clamp(127.5 * image + 128.0, 0, 255).to("cpu", dtype=torch.uint8).numpy() + Image.fromarray(image).save('output.jpg') + print('🤖️: 图片保存至output.jpg') + print('\n') + + + image_dir = './dataset/eval_images/' + prompt = f"一个年轻人准备踢足球" + chat_with_vlm(prompt) \ No newline at end of file diff --git a/model/model_t2i.py b/model/model_t2i.py index 4144406..2af43ef 100644 --- a/model/model_t2i.py +++ b/model/model_t2i.py @@ -3,28 +3,28 @@ from .model import * from typing import Optional, Tuple, List from torch import nn import warnings -from model.model_vlm import MiniMindVLM +from model.model_vlm import MiniMindVLM, VisionProj 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 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), +# nn.ReLU(), +# nn.Linear(self.lm_dim, self.lm_dim) +# ) + +# def forward(self, image_encoders): +# vision_proj = self.vision_proj(image_encoders) +# return vision_proj # 继承自语言模型 @@ -36,6 +36,7 @@ class MiniMindT2I(MiniMindVLM): if not params: params = VLMConfig() self.params = params self.image_tokenizer = self.__class__.get_image_tokenizer() + # self.output = nn.Linear(params.dim, 16384) self.vision_proj = VisionProj(ve_dim=8, lm_dim=params.dim) @staticmethod @@ -116,34 +117,39 @@ class MiniMindT2I(MiniMindVLM): **args): start_pos = args.get('start_pos', 0) pixel_tensors = args.get('pixel_tensors', None) - h = self.tok_embeddings(input_ids) + is_image = args.get('is_image', False) + if is_image: + input_ids = self.image_tokenizer.quantize.get_codebook_entry(input_ids) + h = self.vision_proj(input_ids) + else: + h = self.tok_embeddings(input_ids) - if pixel_tensors is not None and start_pos == 0: - 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]) + if pixel_tensors is not None and start_pos == 0: + 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 = [] @@ -163,3 +169,31 @@ class MiniMindT2I(MiniMindVLM): self.OUT.__setitem__('past_key_values', past_kvs) self.OUT.__setitem__('target_ids', target_ids) # 添加target_ids return self.OUT + + @torch.inference_mode() + def generate(self, input_ids, img_token_num=256, temperature=0.75, top_p=0.90, rp=1., use_cache=True, pad_token_id=0, **args): + idx = 0 + input_ids = input_ids[input_ids != pad_token_id].unsqueeze(0) + start, first_seq, past_kvs = input_ids.shape[1], True, None + while idx < img_token_num: + if first_seq or not use_cache: + out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False + else: + out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache, + start_pos=input_ids.shape[1] - 1, is_image=True, **args) + logits, past_kvs = out.logits[:, -1, :], out.past_key_values + logits[:, list(set(input_ids.tolist()[0]))] /= rp + logits /= (temperature + 1e-9) + if top_p is not None and top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) + sorted_probs = F.softmax(sorted_logits, dim=-1) + cumulative_probs = torch.cumsum(sorted_probs, dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() + sorted_indices_to_remove[:, 0] = False + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = -float('Inf') + input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) + input_ids = torch.cat((input_ids, input_ids_next), dim=1) + idx += 1 + return input_ids[:, -img_token_num:] \ No newline at end of file diff --git a/train_pretrain_t2i.py b/train_pretrain_t2i.py index 70b6190..3e25827 100644 --- a/train_pretrain_t2i.py +++ b/train_pretrain_t2i.py @@ -87,7 +87,7 @@ def train_epoch(epoch, wandb): 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' + ckp = f'{args.save_dir}/pretrain_t2i_{model_config.dim}{moe_path}.pth' if isinstance(model, torch.nn.parallel.DistributedDataParallel): state_dict = model.module.state_dict() else: diff --git a/train_sft_t2i.py b/train_sft_t2i.py new file mode 100644 index 0000000..25fd5df --- /dev/null +++ b/train_sft_t2i.py @@ -0,0 +1,205 @@ +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_t2i import MiniMindT2I +from model.VLMConfig import VLMConfig +from model.dataset import T2IDataset + +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 + 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) + ).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}/sft_t2i_{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/pretrain_t2i_{model_config.dim}{moe_path}.pth' + model = MiniMindT2I(model_config) + state_dict = torch.load(ckp, map_location=args.device) + model.load_state_dict(state_dict, strict=False) + + Logger(f'T2I可训练参数量:{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=6) + 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_t2i_data.jsonl") + 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) + 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) + 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, + image_special_token = '@' * 256, image_ids = [34] * 256) + 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 = 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, + 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)