From b45247a32dc6a10fa487c8c9b42beba804b7db09 Mon Sep 17 00:00:00 2001 From: xinyanghuang7 Date: Wed, 30 Oct 2024 07:42:13 -0400 Subject: [PATCH] multi-images training --- 4-sft_multi_vlm.py | 244 ------------------------------------------- 5-eval_chat_multi.py | 167 ----------------------------- 2 files changed, 411 deletions(-) delete mode 100644 4-sft_multi_vlm.py delete mode 100644 5-eval_chat_multi.py diff --git a/4-sft_multi_vlm.py b/4-sft_multi_vlm.py deleted file mode 100644 index 7ed3d23..0000000 --- a/4-sft_multi_vlm.py +++ /dev/null @@ -1,244 +0,0 @@ -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]) - - 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 deleted file mode 100644 index 59b4c5e..0000000 --- a/5-eval_chat_multi.py +++ /dev/null @@ -1,167 +0,0 @@ -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 - - i = 0 - # 获取图像文件列表 - # image_dir = './dataset/eval_multi_images/' # 对于图像转换数据集 - image_dir = './dataset/eval_multi_images/bird/' # 对于鸟类数据集 - 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?" - - # 对于鸟类数据集 - prompt = "\n\nName all the differences between these two birds." - - # 遍历图像文件 - 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