github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[feat] support parquet data

pull/107/head
jingyaogong 2025-12-27 10:49:50 +08:00
parent 1863743ee9
commit 3684285a91
4 changed files with 37 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
dataset/lm_dataset.py
View File

@ -1,21 +1,24 @@
import sys
import os
__package__ = "dataset"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import json
import torch
import io
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torch
from model.model_vlm import MiniMindVLM
import os
import pyarrow.parquet as pq
os.environ["TOKENIZERS_PARALLELISM"] = "false"
class VLMDataset(Dataset):
def __init__(self, jsonl_path, images_path, tokenizer, preprocess=None, max_length=512,
def __init__(self, parquet_path, tokenizer, preprocess=None, max_length=512,
image_special_token='@' * 196):
super().__init__()
self.samples = self.load_data(jsonl_path)
self.images_path = images_path
self.table = pq.read_table(parquet_path)
self.tokenizer = tokenizer
self.max_length = max_length
self.preprocess = preprocess
@ -24,15 +27,7 @@ class VLMDataset(Dataset):
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
def __len__(self):
return len(self.samples)
def load_data(self, path):
samples = []
with open(path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
data = json.loads(line.strip())
samples.append(data)
return samples
return len(self.table)
def _create_chat_prompt(self, conversations):
messages = []
@ -64,9 +59,10 @@ class VLMDataset(Dataset):
return loss_mask
def __getitem__(self, index: int):
sample = self.samples[index]
image_paths = sample['image']
prompt = self._create_chat_prompt(sample['conversations'])
conversations = json.loads(self.table['conversations'][index].as_py())
image_bytes = self.table['image_bytes'][index].as_py()
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))
loss_mask = self._generate_loss_mask(input_ids)
@ -75,12 +71,16 @@ class VLMDataset(Dataset):
Y = torch.tensor(input_ids[1:], dtype=torch.long)
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
image_tensors = []
for image_name in image_paths.split(','):
image_name = image_name.strip()
image = Image.open(f'{self.images_path}/{image_name}')
image_tensor = MiniMindVLM.image2tensor(image, self.preprocess)
image_tensors.append(image_tensor)
image_tensors = torch.stack(image_tensors, dim=0)
image = Image.open(io.BytesIO(image_bytes))
image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0)
return X, Y, loss_mask, image_tensors
return X, Y, loss_mask, image_tensor
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']:
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')
ax[i].set_title(json.loads(t['conversations'][i].as_py())[1]['content'][:30], fontsize=8)
out = path.replace('.parquet', '_preview.png'); plt.savefig(out); print(f'已保存{out}, 共{len(t)}条')

8
trainer/train_pretrain_vlm.py
View File

@ -84,7 +84,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
del state_dict, clean_state_dict
del X, Y, loss_mask, pixel_values, res, loss
torch.cuda.empty_cache()
if __name__ == "__main__":
@ -100,13 +99,12 @@ if __name__ == "__main__":
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
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=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("--data_path", type=str, default="../dataset/pretrain_data.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）")
@ -142,7 +140,7 @@ if __name__ == "__main__":
# ========== 5. 定义模型、数据、优化器 ==========
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight,
device=args.device, freeze_llm=bool(args.freeze_llm))
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
image_special_token=vlm_config.image_special_token,
max_length=vlm_config.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None

8
trainer/train_sft_vlm.py
View File

@ -84,7 +84,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
del state_dict, clean_state_dict
del X, Y, loss_mask, pixel_values, res, loss
torch.cuda.empty_cache()
if __name__ == "__main__":
@ -100,13 +99,12 @@ if __name__ == "__main__":
parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
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=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("--data_path", type=str, default="../dataset/sft_data.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_wandb", action="store_true", help="是否使用wandb")
@ -141,7 +139,7 @@ if __name__ == "__main__":
# ========== 5. 定义模型、数据、优化器 ==========
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight,
device=args.device)
train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
image_special_token=vlm_config.image_special_token,
max_length=vlm_config.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None

11
trainer/trainer_utils.py
View File

@ -65,11 +65,11 @@ def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../mo
if 'vision_proj' not in name:
param.requires_grad = False
# 默认全参训练时的可选配置（已注释）
# # 只解冻注意力机制中的投影层参数
# for name, param in model.model.named_parameters():
# if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
# param.requires_grad = True
# pretrain阶段：解冻最后1层
last_layer_idx = vlm_config.num_hidden_layers - 1
for name, param in model.model.named_parameters():
if f'layers.{last_layer_idx}.' in name:
param.requires_grad = True
Logger(f'所加载VLM Model可训练参数：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
preprocess = model.processor
@ -121,7 +121,6 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None
torch.save(resume_data, resume_tmp)
os.replace(resume_tmp, resume_path)
del state_dict, clean_state_dict, resume_data
gc.collect()
torch.cuda.empty_cache()
else: # 加载模式
if os.path.exists(resume_path):