#!/usr/bin/env python3 """ Mini BERT预训练模型效果指标测试脚本 适配分布式训练后的模型评估 """ import os import json import time import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from transformers import ( BertForPreTraining, BertTokenizer, BertConfig, DataCollatorForLanguageModeling ) import numpy as np from typing import Dict, List, Tuple, Optional import argparse from tqdm import tqdm import psutil import GPUtil from datetime import datetime import matplotlib.pyplot as plt import seaborn as sns import glob import re class PretrainingEvaluator: """预训练模型评估器 - 适配分布式训练输出""" def __init__(self, model_path: str, tokenizer_path: str, device: str = "cuda"): self.device = device if torch.cuda.is_available() else "cpu" self.model_path = model_path self.tokenizer_path = tokenizer_path # 加载模型和分词器 self.model, self.tokenizer = self._load_model_and_tokenizer() self.model.eval() # 评估指标存储 self.metrics = { 'mlm_accuracy': [], 'mlm_perplexity': [], 'nsp_accuracy': [], 'sop_accuracy': [], 'throughput': [], 'memory_usage': [], 'training_steps': [] } def _find_latest_checkpoint(self, model_path: str) -> str: """找到最新的检查点""" print(f"🔍 查找最新检查点...") # 查找所有step-*文件夹 step_dirs = [] for item in os.listdir(model_path): if item.startswith('step-') and os.path.isdir(os.path.join(model_path, item)): try: step_num = int(item.split('-')[1]) step_dirs.append((step_num, item)) except ValueError: continue if not step_dirs: raise FileNotFoundError(f"未找到任何检查点文件夹在: {model_path}") # 按步数排序，取最新的 step_dirs.sort(key=lambda x: x[0], reverse=True) latest_step, latest_dir = step_dirs[0] latest_path = os.path.join(model_path, latest_dir) print(f"✅ 找到最新检查点: {latest_dir} (步数: {latest_step})") return latest_path def _load_model_and_tokenizer(self): """加载预训练模型和分词器""" print(f"🔄 加载模型从: {self.model_path}") print(f"🔄 加载分词器从: {self.tokenizer_path}") # 检查模型路径是否存在 if not os.path.exists(self.model_path): raise FileNotFoundError(f"模型路径不存在: {self.model_path}") if not os.path.exists(self.tokenizer_path): raise FileNotFoundError(f"分词器路径不存在: {self.tokenizer_path}") # 加载分词器 tokenizer = BertTokenizer.from_pretrained(self.tokenizer_path) # 找到最新检查点 checkpoint_path = self._find_latest_checkpoint(self.model_path) # 检查检查点文件夹中的文件 checkpoint_files = os.listdir(checkpoint_path) print(f"📁 检查点文件: {checkpoint_files}") # 查找模型权重文件 model_files = [] for file in checkpoint_files: if file.endswith(('.bin', '.safetensors', '.pt', '.pth')): model_files.append(file) if not model_files: raise FileNotFoundError(f"在检查点 {checkpoint_path} 中未找到模型权重文件") print(f"📦 找到模型文件: {model_files}") # 尝试加载config.json config_path = os.path.join(checkpoint_path, "config.json") if os.path.exists(config_path): config = BertConfig.from_pretrained(checkpoint_path) print("✅ 从检查点加载配置") else: # 如果没有config.json，使用默认配置 print("⚠️ 未找到config.json，使用默认配置") config = BertConfig( vocab_size=tokenizer.vocab_size, hidden_size=512, num_hidden_layers=6, num_attention_heads=8, intermediate_size=2048, max_position_embeddings=512, type_vocab_size=2, hidden_dropout_prob=0.1, attention_probs_dropout_prob=0.1 ) # 加载预训练模型 try: # 首先尝试从检查点加载 model = BertForPreTraining.from_pretrained(checkpoint_path, config=config) print("✅ 从检查点加载模型成功") # 检查未初始化的权重 uninitialized_params = [] for name, param in model.named_parameters(): if param.requires_grad and param.data.sum() == 0: uninitialized_params.append(name) if uninitialized_params: print("ℹ️ 以下权重未初始化（这是正常的，用于下游任务）:") for param_name in uninitialized_params: print(f" - {param_name}") print(" 这些权重不会影响预训练任务的评估") except Exception as e: print(f"⚠️ 从检查点加载失败: {e}") print("🔄 尝试创建新模型并加载权重...") # 创建新模型 model = BertForPreTraining(config) # 尝试加载权重 weight_loaded = False for model_file in model_files: try: weight_path = os.path.join(checkpoint_path, model_file) if model_file.endswith('.safetensors'): from safetensors.torch import load_file state_dict = load_file(weight_path) else: state_dict = torch.load(weight_path, map_location='cpu') # 处理分布式训练的权重（移除module.前缀） if any(key.startswith('module.') for key in state_dict.keys()): state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()} model.load_state_dict(state_dict, strict=False) weight_loaded = True print(f"✅ 成功加载权重: {model_file}") break except Exception as e2: print(f"⚠️ 加载权重失败 {model_file}: {e2}") continue if not weight_loaded: raise RuntimeError("无法加载任何模型权重文件") model.to(self.device) print(f"✅ 模型加载完成") print(f" - 词汇表大小: {config.vocab_size}") print(f" - 隐藏层维度: {config.hidden_size}") print(f" - 层数: {config.num_hidden_layers}") print(f" - 注意力头数: {config.num_attention_heads}") print(f" - 最大位置编码: {config.max_position_embeddings}") return model, tokenizer def load_training_logs(self, output_dir: str) -> Dict[str, List[float]]: """加载训练日志""" print(f"📊 加载训练日志从: {output_dir}") logs = { 'steps': [], 'train_loss': [], 'val_loss': [], 'learning_rate': [], 'grad_norm': [] } # 查找日志文件 log_files = [] for pattern in ['*.log', '*.txt', 'training_logs*', 'logs*', '*.json']: log_files.extend(glob.glob(os.path.join(output_dir, pattern))) # 也检查各个检查点文件夹 for item in os.listdir(output_dir): if item.startswith('step-'): step_dir = os.path.join(output_dir, item) for pattern in ['*.log', '*.txt', '*.json']: log_files.extend(glob.glob(os.path.join(step_dir, pattern))) if not log_files: print("⚠️ 未找到训练日志文件，跳过收敛曲线绘制") return None print(f"📁 找到日志文件: {log_files}") # 解析日志文件 for log_file in log_files: try: with open(log_file, 'r', encoding='utf-8') as f: for line in f: if 'step' in line.lower() and 'loss' in line.lower(): # 解析日志行，提取步数、损失等信息 # 这里需要根据实际日志格式调整 pass except Exception as e: print(f"⚠️ 解析日志文件失败 {log_file}: {e}") if not logs['steps']: print("⚠️ 无法解析训练日志，跳过收敛曲线绘制") return None return logs def _generate_sample_logs(self) -> Dict[str, List[float]]: """生成示例训练日志数据""" steps = list(range(0, 100000, 1000)) train_loss = [2.5 * np.exp(-s/20000) + 0.5 + 0.1 * np.random.randn() for s in steps] val_loss = [2.8 * np.exp(-s/25000) + 0.6 + 0.15 * np.random.randn() for s in steps] learning_rate = [2e-4 * np.exp(-s/50000) for s in steps] grad_norm = [1.0 * np.exp(-s/15000) + 0.1 + 0.05 * np.random.randn() for s in steps] return { 'steps': steps, 'train_loss': train_loss, 'val_loss': val_loss, 'learning_rate': learning_rate, 'grad_norm': grad_norm } def evaluate_mlm(self, dataloader: DataLoader, num_samples: int = 1000) -> Dict[str, float]: """评估掩码语言建模(MLM)任务""" print("🔍 评估MLM任务...") total_correct = 0 total_tokens = 0 total_loss = 0.0 num_batches = 0 with torch.no_grad(): for batch_idx, batch in enumerate(tqdm(dataloader, desc="MLM评估")): if batch_idx * dataloader.batch_size >= num_samples: break # 移动数据到设备 input_ids = batch['input_ids'].to(self.device) attention_mask = batch['attention_mask'].to(self.device) labels = batch['labels'].to(self.device) # 前向传播 try: outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) except Exception as e: print(f"⚠️ 批次 {batch_idx} 前向传播失败: {e}") continue # 检查输出 if outputs is None: print(f"⚠️ 批次 {batch_idx} 输出为None，跳过") continue # 计算MLM损失和预测 mlm_loss = outputs.loss logits = outputs.prediction_logits # 调试信息 if batch_idx == 0: print(f"🔍 调试信息 - 批次 {batch_idx}:") print(f" - input_ids shape: {input_ids.shape}") print(f" - labels shape: {labels.shape}") print(f" - outputs type: {type(outputs)}") print(f" - mlm_loss: {mlm_loss}") print(f" - logits shape: {logits.shape if logits is not None else 'None'}") # 如果loss为None，手动计算 if mlm_loss is None: print(f"⚠️ 批次 {batch_idx} loss为None，手动计算") # 手动计算MLM损失 loss_fct = nn.CrossEntropyLoss() active_loss = labels.view(-1) != -100 active_logits = logits.view(-1, logits.size(-1)) active_labels = torch.where( active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels) ) mlm_loss = loss_fct(active_logits, active_labels) # 计算准确率 predictions = torch.argmax(logits, dim=-1) mask_positions = (labels != -100) if mask_positions.sum() > 0: correct = (predictions[mask_positions] == labels[mask_positions]).sum().item() total_correct += correct total_tokens += mask_positions.sum().item() if mlm_loss is not None: total_loss += mlm_loss.item() num_batches += 1 # 计算指标 accuracy = total_correct / total_tokens if total_tokens > 0 else 0.0 avg_loss = total_loss / num_batches if num_batches > 0 else float('inf') perplexity = torch.exp(torch.tensor(avg_loss)).item() mlm_metrics = { 'accuracy': accuracy, 'perplexity': perplexity, 'loss': avg_loss, 'total_tokens': total_tokens, 'correct_predictions': total_correct } print(f"✅ MLM评估完成:") print(f" - 准确率: {accuracy:.4f} ({accuracy*100:.2f}%)") print(f" - 困惑度: {perplexity:.2f}") print(f" - 平均损失: {avg_loss:.4f}") print(f" - 总token数: {total_tokens}") print(f" - 正确预测数: {total_correct}") return mlm_metrics def evaluate_nsp(self, dataloader: DataLoader, num_samples: int = 1000) -> Dict[str, float]: """评估下一句预测(NSP)任务""" print("🔍 评估NSP任务...") total_correct = 0 total_samples = 0 with torch.no_grad(): for batch_idx, batch in enumerate(tqdm(dataloader, desc="NSP评估")): if batch_idx * dataloader.batch_size >= num_samples: break # 移动数据到设备 input_ids = batch['input_ids'].to(self.device) attention_mask = batch['attention_mask'].to(self.device) token_type_ids = batch.get('token_type_ids', None) if token_type_ids is not None: token_type_ids = token_type_ids.to(self.device) # 检查是否有NSP标签 if 'next_sentence_label' not in batch: print("⚠️ 数据集中没有NSP标签，跳过NSP评估") return {'accuracy': 0.0, 'total_samples': 0} next_sentence_labels = batch['next_sentence_label'].to(self.device) # 前向传播 outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, next_sentence_label=next_sentence_labels ) # 检查输出 if outputs is None: print(f"⚠️ 批次 {batch_idx} 输出为None，跳过") continue # 计算NSP准确率 seq_relationship_logits = outputs.seq_relationship_logits if seq_relationship_logits is None: print(f"⚠️ 批次 {batch_idx} seq_relationship_logits为None，跳过") continue predictions = torch.argmax(seq_relationship_logits, dim=-1) correct = (predictions == next_sentence_labels).sum().item() total_correct += correct total_samples += next_sentence_labels.size(0) accuracy = total_correct / total_samples if total_samples > 0 else 0.0 nsp_metrics = { 'accuracy': accuracy, 'total_samples': total_samples, 'correct_predictions': total_correct } print(f"✅ NSP评估完成:") print(f" - 准确率: {accuracy:.4f} ({accuracy*100:.2f}%)") print(f" - 总样本数: {total_samples}") print(f" - 正确预测数: {total_correct}") return nsp_metrics def measure_throughput(self, dataloader: DataLoader, num_batches: int = 100) -> Dict[str, float]: """测量模型吞吐量""" print("⚡ 测量模型吞吐量...") batch_times = [] tokens_per_second = [] with torch.no_grad(): for batch_idx, batch in enumerate(tqdm(dataloader, desc="吞吐量测试")): if batch_idx >= num_batches: break # 移动数据到设备 input_ids = batch['input_ids'].to(self.device) attention_mask = batch['attention_mask'].to(self.device) labels = batch['labels'].to(self.device) # 测量推理时间 start_time = time.time() outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) # 同步GPU操作 if self.device == 'cuda': torch.cuda.synchronize() end_time = time.time() batch_time = end_time - start_time batch_times.append(batch_time) # 计算tokens per second batch_size = input_ids.size(0) seq_length = input_ids.size(1) total_tokens = batch_size * seq_length tps = total_tokens / batch_time tokens_per_second.append(tps) avg_batch_time = np.mean(batch_times) avg_tps = np.mean(tokens_per_second) throughput_metrics = { 'avg_batch_time': avg_batch_time, 'avg_tokens_per_second': avg_tps, 'batches_per_second': 1.0 / avg_batch_time, 'batch_times': batch_times, 'tokens_per_second': tokens_per_second } print(f"✅ 吞吐量测量完成:") print(f" - 平均批次时间: {avg_batch_time:.4f}s") print(f" - 平均tokens/秒: {avg_tps:.2f}") print(f" - 批次/秒: {1.0/avg_batch_time:.2f}") return throughput_metrics def measure_memory_usage(self) -> Dict[str, float]: """测量内存使用情况""" print("💾 测量内存使用...") memory_metrics = {} # CPU内存使用 cpu_memory = psutil.virtual_memory() memory_metrics['cpu_total_gb'] = cpu_memory.total / (1024**3) memory_metrics['cpu_used_gb'] = cpu_memory.used / (1024**3) memory_metrics['cpu_available_gb'] = cpu_memory.available / (1024**3) memory_metrics['cpu_percent'] = cpu_memory.percent # GPU内存使用 if torch.cuda.is_available(): gpu_memory = torch.cuda.memory_allocated() / (1024**3) gpu_max_memory = torch.cuda.max_memory_allocated() / (1024**3) gpu_reserved = torch.cuda.memory_reserved() / (1024**3) memory_metrics['gpu_allocated_gb'] = gpu_memory memory_metrics['gpu_max_allocated_gb'] = gpu_max_memory memory_metrics['gpu_reserved_gb'] = gpu_reserved # 获取GPU信息 try: gpus = GPUtil.getGPUs() if gpus: gpu = gpus[0] memory_metrics['gpu_total_gb'] = gpu.memoryTotal / 1024 memory_metrics['gpu_used_gb'] = gpu.memoryUsed / 1024 memory_metrics['gpu_free_gb'] = gpu.memoryFree / 1024 memory_metrics['gpu_utilization'] = gpu.load * 100 except Exception as e: print(f"⚠️ 获取GPU信息失败: {e}") print(f"✅ 内存使用测量完成:") print(f" - CPU使用率: {memory_metrics['cpu_percent']:.1f}%") print(f" - CPU已用: {memory_metrics['cpu_used_gb']:.2f}GB") if 'gpu_allocated_gb' in memory_metrics: print(f" - GPU已分配: {memory_metrics['gpu_allocated_gb']:.2f}GB") print(f" - GPU最大分配: {memory_metrics['gpu_max_allocated_gb']:.2f}GB") return memory_metrics def plot_convergence_curves(self, training_logs: Dict[str, List[float]]): """绘制收敛曲线""" print("📊 绘制收敛曲线...") steps = training_logs['steps'] train_loss = training_logs['train_loss'] val_loss = training_logs['val_loss'] learning_rate = training_logs['learning_rate'] grad_norm = training_logs['grad_norm'] # 创建图表 plt.figure(figsize=(15, 10)) # 子图1: 损失曲线 plt.subplot(2, 3, 1) plt.plot(steps, train_loss, label='训练损失', color='blue', alpha=0.7) plt.plot(steps, val_loss, label='验证损失', color='red', alpha=0.7) plt.xlabel('训练步数') plt.ylabel('损失值') plt.title('训练收敛曲线') plt.legend() plt.grid(True, alpha=0.3) plt.yscale('log') # 子图2: 学习率曲线 plt.subplot(2, 3, 2) plt.plot(steps, learning_rate, label='学习率', color='green') plt.xlabel('训练步数') plt.ylabel('学习率') plt.title('学习率调度') plt.legend() plt.grid(True, alpha=0.3) plt.yscale('log') # 子图3: 梯度范数 plt.subplot(2, 3, 3) plt.plot(steps, grad_norm, label='梯度范数', color='brown') plt.xlabel('训练步数') plt.ylabel('梯度范数') plt.title('梯度范数变化') plt.legend() plt.grid(True, alpha=0.3) # 子图4: 损失对比（线性尺度） plt.subplot(2, 3, 4) plt.plot(steps, train_loss, label='训练损失', color='blue', alpha=0.7) plt.plot(steps, val_loss, label='验证损失', color='red', alpha=0.7) plt.xlabel('训练步数') plt.ylabel('损失值') plt.title('损失对比（线性尺度）') plt.legend() plt.grid(True, alpha=0.3) # 子图5: 学习率与损失关系 plt.subplot(2, 3, 5) plt.scatter(learning_rate, train_loss, alpha=0.6, s=1) plt.xlabel('学习率') plt.ylabel('训练损失') plt.title('学习率与损失关系') plt.grid(True, alpha=0.3) plt.xscale('log') plt.yscale('log') # 子图6: 梯度范数与损失关系 plt.subplot(2, 3, 6) plt.scatter(grad_norm, train_loss, alpha=0.6, s=1) plt.xlabel('梯度范数') plt.ylabel('训练损失') plt.title('梯度范数与损失关系') plt.grid(True, alpha=0.3) plt.yscale('log') plt.tight_layout() plt.savefig('pretraining_convergence_curves.png', dpi=300, bbox_inches='tight') plt.show() print("✅ 收敛曲线已保存到: pretraining_convergence_curves.png") def run_comprehensive_evaluation(self, dataloader: DataLoader, num_samples: int = 1000, training_logs: Dict[str, List[float]] = None) -> Dict[str, any]: """运行综合评估""" print("🚀 开始综合评估...") print("="*60) evaluation_results = {} # 1. MLM评估 mlm_results = self.evaluate_mlm(dataloader, num_samples) evaluation_results['mlm'] = mlm_results # 2. NSP评估 nsp_results = self.evaluate_nsp(dataloader, num_samples) evaluation_results['nsp'] = nsp_results # 3. 吞吐量测量 throughput_results = self.measure_throughput(dataloader, min(100, num_samples//10)) evaluation_results['throughput'] = throughput_results # 4. 内存使用测量 memory_results = self.measure_memory_usage() evaluation_results['memory'] = memory_results # 5. 绘制收敛曲线 if training_logs: self.plot_convergence_curves(training_logs) else: print("⚠️ 未提供训练日志，跳过收敛曲线绘制") # 6. 生成评估报告 self._generate_evaluation_report(evaluation_results, training_logs) print("="*60) print("✅ 综合评估完成!") return evaluation_results def _generate_evaluation_report(self, results: Dict[str, any], training_logs: Dict[str, List[float]] = None): """生成评估报告""" report = { 'evaluation_time': datetime.now().isoformat(), 'model_path': self.model_path, 'tokenizer_path': self.tokenizer_path, 'device': self.device, 'results': results, 'training_logs': training_logs } # 保存JSON报告 with open('pretraining_evaluation_report.json', 'w', encoding='utf-8') as f: json.dump(report, f, ensure_ascii=False, indent=2) # 生成文本报告 with open('pretraining_evaluation_summary.txt', 'w', encoding='utf-8') as f: f.write("Mini BERT预训练模型评估报告\n") f.write("="*50 + "\n\n") f.write(f"评估时间: {report['evaluation_time']}\n") f.write(f"模型路径: {report['model_path']}\n") f.write(f"分词器路径: {report['tokenizer_path']}\n") f.write(f"运行设备: {report['device']}\n\n") # MLM结果 mlm = results['mlm'] f.write("掩码语言建模(MLM)结果:\n") f.write(f" 准确率: {mlm['accuracy']:.4f} ({mlm['accuracy']*100:.2f}%)\n") f.write(f" 困惑度: {mlm['perplexity']:.2f}\n") f.write(f" 平均损失: {mlm['loss']:.4f}\n") f.write(f" 总token数: {mlm['total_tokens']}\n") f.write(f" 正确预测数: {mlm['correct_predictions']}\n\n") # NSP结果 nsp = results['nsp'] f.write("下一句预测(NSP)结果:\n") f.write(f" 准确率: {nsp['accuracy']:.4f} ({nsp['accuracy']*100:.2f}%)\n") f.write(f" 总样本数: {nsp['total_samples']}\n") f.write(f" 正确预测数: {nsp['correct_predictions']}\n\n") # 吞吐量结果 throughput = results['throughput'] f.write("性能指标:\n") f.write(f" 平均批次时间: {throughput['avg_batch_time']:.4f}s\n") f.write(f" 平均tokens/秒: {throughput['avg_tokens_per_second']:.2f}\n") f.write(f" 批次/秒: {1.0/throughput['avg_batch_time']:.2f}\n\n") # 内存使用 memory = results['memory'] f.write("内存使用:\n") f.write(f" CPU使用率: {memory['cpu_percent']:.1f}%\n") f.write(f" CPU已用: {memory['cpu_used_gb']:.2f}GB\n") if 'gpu_allocated_gb' in memory: f.write(f" GPU已分配: {memory['gpu_allocated_gb']:.2f}GB\n") f.write(f" GPU最大分配: {memory['gpu_max_allocated_gb']:.2f}GB\n") # 训练日志摘要 if training_logs: f.write("\n训练日志摘要:\n") f.write(f" 总训练步数: {len(training_logs['steps'])}\n") f.write(f" 最终训练损失: {training_logs['train_loss'][-1]:.4f}\n") f.write(f" 最终验证损失: {training_logs['val_loss'][-1]:.4f}\n") f.write(f" 最终学习率: {training_logs['learning_rate'][-1]:.2e}\n") f.write(f" 最终梯度范数: {training_logs['grad_norm'][-1]:.4f}\n") print("📄 评估报告已生成:") print(" - pretraining_evaluation_report.json (详细JSON报告)") print(" - pretraining_evaluation_summary.txt (文本摘要)") if training_logs: print(" - pretraining_convergence_curves.png (收敛曲线图)") def create_real_dataloader(data_path: str, tokenizer_path: str, batch_size: int = 8, max_length: int = 128): """创建真实预训练数据加载器""" class RealPretrainingDataset(Dataset): def __init__(self, data_path: str, tokenizer, max_length: int = 128): self.tokenizer = tokenizer self.max_length = max_length self.samples = [] print(f"📖 加载真实预训练数据: {data_path}") # 读取JSONL数据 with open(data_path, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f): try: sample = json.loads(line.strip()) self.samples.append(sample) except json.JSONDecodeError as e: print(f"⚠️ 跳过无效JSON行 {line_num + 1}: {e}") continue print(f"✅ 加载了 {len(self.samples)} 个真实样本") # 显示样本结构 if self.samples: print(f"🔍 样本结构示例:") sample_keys = list(self.samples[0].keys()) print(f" - 字段: {sample_keys}") for key in sample_keys: if isinstance(self.samples[0][key], (list, torch.Tensor)): print(f" - {key} shape: {len(self.samples[0][key]) if isinstance(self.samples[0][key], list) else self.samples[0][key].shape}") else: print(f" - {key}: {type(self.samples[0][key])}") def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] # 转换为tensor input_ids = torch.tensor(sample['input_ids'], dtype=torch.long) attention_mask = torch.tensor(sample['attention_mask'], dtype=torch.long) labels = torch.tensor(sample['labels'], dtype=torch.long) # 处理token_type_ids if 'token_type_ids' in sample: token_type_ids = torch.tensor(sample['token_type_ids'], dtype=torch.long) else: token_type_ids = torch.zeros_like(input_ids) # 处理NSP标签 if 'next_sentence_label' in sample: next_sentence_label = torch.tensor(sample['next_sentence_label'], dtype=torch.long) else: next_sentence_label = torch.tensor(0, dtype=torch.long) return { 'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids, 'labels': labels, 'next_sentence_label': next_sentence_label } # 加载分词器 tokenizer = BertTokenizer.from_pretrained(tokenizer_path) # 创建数据集 dataset = RealPretrainingDataset(data_path, tokenizer, max_length) # 创建数据加载器 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) return dataloader def create_sample_dataloader(tokenizer_path: str, batch_size: int = 8, max_length: int = 128): """创建示例数据加载器用于测试""" class SampleDataset(Dataset): def __init__(self, tokenizer, num_samples=1000, max_length=128): self.tokenizer = tokenizer self.num_samples = num_samples self.max_length = max_length self.special_tokens = { 'cls': tokenizer.cls_token_id, 'sep': tokenizer.sep_token_id, 'pad': tokenizer.pad_token_id, 'mask': tokenizer.mask_token_id } def __len__(self): return self.num_samples def __getitem__(self, idx): # 生成示例文本 sample_texts = [ "这是一个测试句子，用于评估预训练模型的效果。", "机器学习是人工智能的一个重要分支。", "深度学习模型需要大量的训练数据。", "自然语言处理技术正在快速发展。", "Transformer架构改变了NLP领域。" ] # 随机选择文本 text = np.random.choice(sample_texts) # 分词 tokens = self.tokenizer.encode(text, add_special_tokens=True, max_length=self.max_length, truncation=True, padding='max_length') input_ids = torch.tensor(tokens, dtype=torch.long) # 创建注意力掩码 attention_mask = (input_ids != self.special_tokens['pad']).long() # 创建MLM标签（随机掩码15%的token） labels = input_ids.clone() mask_positions = torch.rand(len(tokens)) < 0.15 labels[~mask_positions] = -100 # 不计算损失的token labels[mask_positions] = input_ids[mask_positions] # 被掩码的token # 随机掩码 input_ids[mask_positions] = self.special_tokens['mask'] # 创建token类型ID（简化处理） token_type_ids = torch.zeros_like(input_ids) # 创建NSP标签（随机） next_sentence_label = torch.tensor(np.random.randint(0, 2), dtype=torch.long) return { 'input_ids': input_ids, 'attention_mask': attention_mask, 'token_type_ids': token_type_ids, 'labels': labels, 'next_sentence_label': next_sentence_label } # 加载分词器 tokenizer = BertTokenizer.from_pretrained(tokenizer_path) # 创建数据集 dataset = SampleDataset(tokenizer, num_samples=1000, max_length=max_length) # 创建数据加载器 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True) return dataloader def main(): parser = argparse.ArgumentParser(description="Mini BERT预训练模型效果评估") parser.add_argument("--model_path", required=True, help="预训练模型路径") parser.add_argument("--tokenizer_path", required=True, help="分词器路径") parser.add_argument("--data_path", required=True, help="评估数据路径（JSONL格式的预训练数据）") parser.add_argument("--num_samples", type=int, default=1000, help="评估样本数量") parser.add_argument("--batch_size", type=int, default=8, help="批次大小") parser.add_argument("--max_length", type=int, default=128, help="最大序列长度") parser.add_argument("--device", default="cuda", help="运行设备") parser.add_argument("--training_logs", help="训练日志文件路径（可选）") args = parser.parse_args() # 检查模型和分词器路径 if not os.path.exists(args.model_path): print(f"❌ 模型路径不存在: {args.model_path}") return if not os.path.exists(args.tokenizer_path): print(f"❌ 分词器路径不存在: {args.tokenizer_path}") return # 创建评估器 evaluator = PretrainingEvaluator( model_path=args.model_path, tokenizer_path=args.tokenizer_path, device=args.device ) # 加载训练日志 training_logs = None if args.training_logs and os.path.exists(args.training_logs): training_logs = evaluator.load_training_logs(args.training_logs) elif os.path.exists(args.model_path): # 尝试从模型输出目录加载日志 training_logs = evaluator.load_training_logs(args.model_path) # 创建数据加载器 - 必须使用真实数据 if not args.data_path: print("❌ 必须提供真实评估数据路径 --data_path") print("💡 请使用以下命令构造评估数据:") print("python create_pretraining_validation_data.py \\") print(" --validation_file /root/autodl-tmp/pretraining_data/validation.txt \\") print(" --tokenizer_path /root/autodl-tmp/pretraining_data/tokenizer \\") print(" --output_file /root/autodl-tmp/pretraining_data/validation.jsonl") return if not os.path.exists(args.data_path): print(f"❌ 评估数据文件不存在: {args.data_path}") print("💡 请先构造评估数据文件") return # 使用真实数据 print(f"📁 使用真实评估数据: {args.data_path}") dataloader = create_real_dataloader(args.data_path, args.tokenizer_path, args.batch_size, args.max_length) # 运行评估 results = evaluator.run_comprehensive_evaluation( dataloader=dataloader, num_samples=args.num_samples, training_logs=training_logs ) print("\n🎉 评估完成！") print("📊 主要指标:") print(f" MLM准确率: {results['mlm']['accuracy']:.4f}") print(f" MLM困惑度: {results['mlm']['perplexity']:.2f}") print(f" NSP准确率: {results['nsp']['accuracy']:.4f}") print(f" 吞吐量: {results['throughput']['avg_tokens_per_second']:.2f} tokens/s") if __name__ == "__main__": main()