From ad1aa51ffdba8a03b1960c3891a38b6dd4d7d843 Mon Sep 17 00:00:00 2001 From: xinyanghuang7 Date: Tue, 29 Oct 2024 10:12:35 -0400 Subject: [PATCH] add multi-img --- 4-sft_multi_vlm.py | 244 +++++++++++++++++++++ 5-eval_chat_multi.py | 164 ++++++++++++++ model/__pycache__/LMConfig.cpython-310.pyc | Bin 1601 -> 1601 bytes 3 files changed, 408 insertions(+) create mode 100644 4-sft_multi_vlm.py create mode 100644 5-eval_chat_multi.py diff --git a/4-sft_multi_vlm.py b/4-sft_multi_vlm.py new file mode 100644 index 0000000..3c820e2 --- /dev/null +++ b/4-sft_multi_vlm.py @@ -0,0 +1,244 @@ +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 +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data import DataLoader, DistributedSampler +from transformers import AutoTokenizer, AutoModel +from model.model import Transformer +from model.LMConfig import LMConfig +from model.dataset import SFTDataset_multi +from model.vision_utils import get_vision_model, get_img_embedding + +warnings.filterwarnings('ignore') + + +def count_parameters(model): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def Logger(content): + if not ddp or dist.get_rank() == 0: + print(content) + + +def get_lr(it, all): + warmup_iters = args.warmup_iters + lr_decay_iters = all + min_lr = args.learning_rate / 10 + + if it < warmup_iters: + return args.learning_rate * it / warmup_iters + if it > lr_decay_iters: + return min_lr + decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters) + assert 0 <= decay_ratio <= 1 + coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) + return min_lr + coeff * (args.learning_rate - min_lr) + + +def train_epoch(epoch, wandb): + start_time = time.time() + for step, (X, Y, loss_mask, image_process) in enumerate(train_loader): + X = X.to(args.device) + Y = Y.to(args.device) + loss_mask = loss_mask.to(args.device) + image_process = image_process.to(args.device) + image_encoders = get_img_embedding(image_process, vision_model) # torch.Size([32, 4, 50, 768]) + lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch) + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + with ctx: + logits = model(X, Y, image_encoders=image_encoders).logits + loss = F.cross_entropy(logits.view(-1, logits.size(-1)), Y.view(-1), ignore_index=0, reduction='none') + loss_mask = loss_mask.view(-1) + loss = torch.sum(loss * loss_mask) / loss_mask.sum() + + 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 lm_config.use_moe else '' + ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth' + + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + + torch.save(state_dict, ckp) + model.train() + + if args.save_last and (not ddp or dist.get_rank() == 0): + model.eval() + moe_path = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth' + + if isinstance(model, torch.nn.parallel.DistributedDataParallel): + state_dict = model.module.state_dict() + else: + state_dict = model.state_dict() + + torch.save(state_dict, ckp) + model.train() + + +def init_model(lm_config): + tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') + moe_path = '_moe' if lm_config.use_moe else '' + # ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth' + ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth' + + model = Transformer(lm_config) + state_dict = torch.load(ckp, map_location=args.device) + + # 处理不需要的前缀 + unwanted_prefix = '_orig_mod.' + for k, v in list(state_dict.items()): + if k.startswith(unwanted_prefix): + state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) + + # 加载到模型中 + model.load_state_dict(state_dict, strict=False) + model = model.to(args.device) + + print(f'模型可学习参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)') + + (vision_model, preprocess) = get_vision_model() + vision_model = vision_model.to(args.device) + return model, tokenizer, (vision_model, preprocess) + + +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-SFT") + parser.add_argument("--out_dir", type=str, default="out", help="Output directory") + parser.add_argument("--epochs", type=int, default=30, help="Number of epochs") + parser.add_argument("--batch_size", type=int, default=32, help="Batch size") + parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate") + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", + help="Device to use") + parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type") + parser.add_argument("--use_wandb", default=False, action="store_true", help="Use Weights & Biases") + parser.add_argument("--wandb_project", type=str, default="MiniMind-V", help="Weights & Biases project name") + parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading") + parser.add_argument("--data_path", type=str, default="./dataset/sft_multi_images/output.json", + help="Path to training data") + parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel") + parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps") + parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold") + parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations") + parser.add_argument("--log_interval", type=int, default=100, help="Logging interval") + parser.add_argument("--save_interval", type=int, default=100, help="Model saving interval") + parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training') + parser.add_argument('--save_last', type=bool, default=True, help='save last step model') + + args = parser.parse_args() + + lm_config = LMConfig() + max_seq_len = lm_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-SFT-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, (vision_model, preprocess) = init_model(lm_config) + + use_version = 0 + train_ds = SFTDataset_multi(args.data_path, tokenizer, vision_model=(vision_model, preprocess), + image_special_token=lm_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.Adam(model.parameters(), lr=args.learning_rate) + + if False and not lm_config.use_moe and platform.system() != 'Windows' and float( + torch.__version__.split('.')[0]) >= 2: + Logger("compiling the model... (takes a ~minute)") + unoptimized_model = model + model = torch.compile(model) + + if ddp: + model._ddp_params_and_buffers_to_ignore = {"pos_cis"} + model = DistributedDataParallel(model, device_ids=[ddp_local_rank], find_unused_parameters=True) + + iter_per_epoch = len(train_loader) + for epoch in range(args.epochs): + train_epoch(epoch, wandb) diff --git a/5-eval_chat_multi.py b/5-eval_chat_multi.py new file mode 100644 index 0000000..fdf4d86 --- /dev/null +++ b/5-eval_chat_multi.py @@ -0,0 +1,164 @@ +import os +import random +import time + +import numpy as np +import torch +import warnings + +from PIL import Image +from transformers import AutoTokenizer, AutoModelForCausalLM +from model.model import Transformer +from model.LMConfig import LMConfig +from model.vision_utils import get_vision_model, get_img_process, get_img_embedding + +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): + tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') + model_from = 1 # 1从权重，2用transformers + + if model_from == 1: + moe_path = '_moe' if lm_config.use_moe else '' + ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft_multi.pth' + + model = Transformer(lm_config) + state_dict = torch.load(ckp, map_location=device) + + # 处理不需要的前缀 + unwanted_prefix = '_orig_mod.' + for k, v in list(state_dict.items()): + if k.startswith(unwanted_prefix): + state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) + + # 加载到模型中 + model.load_state_dict(state_dict, strict=False) + else: + model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small', + trust_remote_code=True) + + model = model.to(device) + + print(f'模型参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)') + + (vision_model, preprocess) = get_vision_model() + vision_model = vision_model.to(device) + return model, tokenizer, (vision_model, preprocess) + + +def setup_seed(seed): + random.seed(seed) # 设置 Python 的随机种子 + np.random.seed(seed) # 设置 NumPy 的随机种子 + torch.manual_seed(seed) # 设置 PyTorch 的随机种子 + torch.cuda.manual_seed(seed) # 为当前 GPU 设置随机种子（如果有） + torch.cuda.manual_seed_all(seed) # 为所有 GPU 设置随机种子（如果有） + torch.backends.cudnn.deterministic = True # 确保每次返回的卷积算法是确定的 + torch.backends.cudnn.benchmark = False # 关闭 cuDNN 的自动调优，避免不确定性 + + +if __name__ == "__main__": + # ----------------------------------------------------------------------------- + out_dir = 'out' + start = "" + temperature = 0.5 + top_k = 8 + setup_seed(1337) + # device = 'cpu' + device = 'cuda:0' if torch.cuda.is_available() else 'cpu' + dtype = 'bfloat16' + max_seq_len = 1 * 1024 + lm_config = LMConfig() + lm_config.max_seq_len = max_seq_len + # ----------------------------------------------------------------------------- + + model, tokenizer, (vision_model, preprocess) = init_model(lm_config) + + model = model.eval() + # 推送到huggingface + # model.push_to_hub("minimind") + # tokenizer.push_to_hub("minimind") + + stream = True + + # prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？' + + i = 0 + # 获取图像文件列表 + image_dir = './dataset/eval_multi_images/' + image_files = sorted(os.listdir(image_dir)) + + prompt = "context: Source Image: Target Image: Instruction: What is the correct image edit instruction that can transfrom the source image to target image?" + + # 遍历图像文件 + image_encoders = [] + for i, image_file in enumerate(image_files): + image_path = os.path.join(image_dir, image_file) + + # 打开图像并转换为RGB格式 + image = Image.open(image_path).convert('RGB') + image_process = get_img_process(image, preprocess).to(vision_model.device) + + # 对图像进行编码 + image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0) + image_encoders.append(image_encoder) + print(f'[Image]: {image_file}') + image_encoders = torch.cat(image_encoders, dim=0).unsqueeze(0) + + # 准备问题和消息 + prompt_ = prompt.replace("\n", "") + print(f'[Q]: {prompt_}') + + messages = [{"role": "user", "content": prompt}] + + new_prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True + )[-(max_seq_len - 1):] + + x = tokenizer(new_prompt).data['input_ids'] + x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...]) + + with torch.no_grad(): + res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature, + top_k=top_k, stream=stream, image_encoders=image_encoders) + + print('[A]: ', end='') + try: + y = next(res_y) + except StopIteration: + print("No answer") + exit() + + history_idx = 0 + while y is not None: + answer = tokenizer.decode(y[0].tolist()) + if answer and answer[-1] == '�': + try: + y = next(res_y) + except StopIteration: + break + continue + if not len(answer): + try: + y = next(res_y) + except StopIteration: + break + continue + + print(answer[history_idx:], end='', flush=True) + try: + y = next(res_y) + except StopIteration: + break + history_idx = len(answer) + + if not stream: + break + + print('\n') \ No newline at end of file diff --git a/model/__pycache__/LMConfig.cpython-310.pyc b/model/__pycache__/LMConfig.cpython-310.pyc index a946bdef8e0e059fa6e56071cec1c58e5f1c6c3c..8f08b709310eb725b4889ddc59842de6fbb4ea3b 100644 GIT binary patch delta 19 ZcmX@ebC8ECpO=@50SKBNY~(U$0{}1V1d;#% delta 19 ZcmX@ebC8ECpO=@50SNT(Z{#v(0{|`c1U>)&