From 7605850e1d3ccdd002f13cf877f4ccc1dd1dd8b7 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 30 Oct 2025 10:50:54 +0800 Subject: [PATCH] [feat] update args --- eval_vlm.py | 6 +++--- scripts/web_demo_vlm.py | 35 ++++++++++++++++------------------- trainer/train_pretrain_vlm.py | 10 +++++----- trainer/train_sft_vlm.py | 6 +++--- 4 files changed, 27 insertions(+), 30 deletions(-) diff --git a/eval_vlm.py b/eval_vlm.py index f4c38d5..93357ce 100644 --- a/eval_vlm.py +++ b/eval_vlm.py @@ -10,11 +10,11 @@ warnings.filterwarnings('ignore') def init_model(args): tokenizer = AutoTokenizer.from_pretrained(args.load_from) - if args.load_from == 'model': + if 'model' in args.load_from: moe_suffix = '_moe' if args.use_moe else '' ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' model = MiniMindVLM( - VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=args.use_moe), + VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)), vision_model_path="./model/vision_model/clip-vit-base-patch16" ) state_dict = torch.load(ckp, map_location=args.device) @@ -35,7 +35,7 @@ def main(): parser.add_argument('--weight', default='sft_vlm', type=str, help="权重名称前缀（pretrain_vlm, sft_vlm）") parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度（512=Small-26M, 768=Base-104M）") parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量（Small=8, Base=16）") - parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE架构") + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") parser.add_argument('--max_new_tokens', default=512, type=int, help="最大生成长度") parser.add_argument('--temperature', default=0.65, type=float, help="生成温度，控制随机性（0-1，越大越随机）") parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）") diff --git a/scripts/web_demo_vlm.py b/scripts/web_demo_vlm.py index 3d0fce2..db867a4 100644 --- a/scripts/web_demo_vlm.py +++ b/scripts/web_demo_vlm.py @@ -19,18 +19,15 @@ warnings.filterwarnings('ignore') def init_model(lm_config): - tokenizer = AutoTokenizer.from_pretrained('../model') - - if args.load == 0: + tokenizer = AutoTokenizer.from_pretrained(args.load_from) + if 'model' in args.load_from: moe_path = '_moe' if lm_config.use_moe else '' - ckp = f'../out/sft_vlm_{lm_config.hidden_size}{moe_path}.pth' + ckp = f'../{args.save_dir}/{args.weight}_{lm_config.hidden_size}{moe_path}.pth' model = MiniMindVLM(lm_config, vision_model_path="../model/vision_model/clip-vit-base-patch16") state_dict = torch.load(ckp, map_location=args.device) model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False) else: - transformers_model_path = '../MiniMind2-V' - tokenizer = AutoTokenizer.from_pretrained(transformers_model_path) - model = AutoModelForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("../model/vision_model/clip-vit-base-patch16") print(f'VLM参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') @@ -182,20 +179,20 @@ def launch_gradio_server(server_name="0.0.0.0", server_port=7788): 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) - parser.add_argument('--hidden_size', default=512, type=int) - parser.add_argument('--num_hidden_layers', default=8, type=int) - parser.add_argument('--max_seq_len', default=8192, type=int) - parser.add_argument('--use_moe', default=False, type=bool) - parser.add_argument('--stream', default=True, type=bool) - parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载") + parser.add_argument('--load_from', default='../model', type=str, help="模型加载路径（model=原生torch权重，其他路径=transformers格式）") + parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录") + parser.add_argument('--weight', default='sft_vlm', type=str, help="权重名称前缀（pretrain_vlm, sft_vlm）") + parser.add_argument('--temperature', default=0.65, type=float, help="生成温度，控制随机性（0-1，越大越随机）") + parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）") + parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度（512=Small-26M, 768=Base-104M）") + parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量（Small=8, Base=16）") + parser.add_argument('--max_seq_len', default=8192, type=int, help="最大序列长度") + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") + parser.add_argument('--stream', default=1, type=int, choices=[0, 1], help="是否使用流式输出（0=否，1=是）") args = parser.parse_args() lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - max_seq_len=args.max_seq_len, use_moe=args.use_moe) + max_seq_len=args.max_seq_len, use_moe=bool(args.use_moe)) model, tokenizer, vision_model, preprocess = init_model(lm_config) launch_gradio_server(server_name="0.0.0.0", server_port=8888) diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index 324cddc..32b155d 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -100,12 +100,12 @@ if __name__ == "__main__": parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度") parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") parser.add_argument('--max_seq_len', default=640, type=int, help="训练的最大截断长度") - parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE") + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl", help="训练数据路径") parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images", help="训练图像路径") parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") - parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是") - parser.add_argument('--freeze_llm', default=True, type=bool, help="是否冻结LLM参数（仅训练vision_proj）") + parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") + parser.add_argument('--freeze_llm', default=1, type=int, choices=[0, 1], help="是否冻结LLM参数（0=否，1=是，仅训练vision_proj）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名") args = parser.parse_args() @@ -118,7 +118,7 @@ if __name__ == "__main__": # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - max_seq_len=args.max_seq_len, use_moe=args.use_moe) + max_seq_len=args.max_seq_len, use_moe=bool(args.use_moe)) ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== @@ -137,7 +137,7 @@ if __name__ == "__main__": # ========== 5. 定义模型、数据、优化器 ========== model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, - device=args.device, freeze_llm=args.freeze_llm) + device=args.device, freeze_llm=bool(args.freeze_llm)) train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess, image_special_token=vlm_config.image_special_token, max_length=vlm_config.max_seq_len) diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 183a43f..253436e 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -100,11 +100,11 @@ if __name__ == "__main__": parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度") parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量") parser.add_argument('--max_seq_len', default=1536, type=int, help="训练的最大截断长度") - parser.add_argument('--use_moe', default=False, type=bool, help="是否使用MoE") + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl", help="训练数据路径") parser.add_argument("--images_path", type=str, default="../dataset/sft_images", help="训练图像路径") parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") - parser.add_argument('--from_resume', default=0, type=int, help="是否自动检测&续训，0否1是") + parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名") args = parser.parse_args() @@ -117,7 +117,7 @@ if __name__ == "__main__": # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) vlm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, - max_seq_len=args.max_seq_len, use_moe=args.use_moe) + max_seq_len=args.max_seq_len, use_moe=bool(args.use_moe)) ckp_data = vlm_checkpoint(vlm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ==========