From 33db76fe102eb58a597503c14cd2b23d46ac4829 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 5 Feb 2026 00:22:58 +0800 Subject: [PATCH] [update] dataset name --- dataset/lm_dataset.py | 6 +++--- trainer/train_pretrain_vlm.py | 2 +- trainer/train_sft_vlm.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index 911de85..5f16559 100644 --- a/dataset/lm_dataset.py +++ b/dataset/lm_dataset.py @@ -61,14 +61,14 @@ class VLMDataset(Dataset): def __getitem__(self, index: int): conversations = json.loads(self.table['conversations'][index].as_py()) image_bytes = self.table['image_bytes'][index].as_py() + if not isinstance(image_bytes, list): image_bytes = [image_bytes] prompt = self.create_chat_prompt(conversations) input_ids = self.tokenizer(prompt).input_ids[:self.max_length] input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) labels = self.generate_labels(input_ids) - image = Image.open(io.BytesIO(image_bytes)) - image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0) + image_tensor = torch.stack([MiniMindVLM.image2tensor(Image.open(io.BytesIO(img)), self.preprocess) for img in image_bytes]) # # === 调试打印 === # print(f"\n--- Sample {index} ---") # for i, (x, y) in enumerate(zip(input_ids[:-1], labels[1:])): @@ -80,7 +80,7 @@ class VLMDataset(Dataset): # 测试parquet数据读取和可视化 if __name__ == '__main__': import matplotlib.pyplot as plt; plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei'] - for path in ['pretrain_data.parquet', 'sft_data.parquet']: + for path in ['pretrain_i2t.parquet', 'sft_i2t.parquet']: t = pq.read_table(path); fig, ax = plt.subplots(1, 5, figsize=(20, 4)) for i in range(5): ax[i].imshow(Image.open(io.BytesIO(t['image_bytes'][i].as_py()))); ax[i].axis('off') diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index 6d29bba..6e1888b 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -95,7 +95,7 @@ if __name__ == "__main__": 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=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") - parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.parquet", help="训练数据路径") + parser.add_argument("--data_path", type=str, default="../dataset/pretrain_i2t.parquet", help="训练数据路径") parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") 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）") diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 8feaba2..75bb8ea 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -95,7 +95,7 @@ if __name__ == "__main__": 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=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）") - parser.add_argument("--data_path", type=str, default="../dataset/sft_data.parquet", help="训练数据路径") + parser.add_argument("--data_path", type=str, default="../dataset/sft_i2t.parquet", help="训练数据路径") parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练") parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")