#!/opt/homebrew/bin/python3 """ 对比 PyTorch 原始模型和 MLX 实现 逐层对比输出，找出差异 目标：找到准确度低的根本原因 """ import torch import mlx.core as mx import numpy as np from PIL import Image, ImageDraw from transformers import AutoModel, AutoProcessor import sys sys.path.append('/Users/gt/.gemini/antigravity/scratch/paddleocr-mlx-conversion') from final_fixed import FinalOptimizedPaddleOCRMLX def create_test_image(): """创建简单的测试图像""" img = Image.new('RGB', (400, 200), color='white') draw = ImageDraw.Draw(img) draw.text((50, 80), "Hello World", fill='black') test_path = "/tmp/test_compare.png" img.save(test_path) return test_path, img def test_pytorch_model(image_path, img): """测试 PyTorch 模型""" print("\n" + "="*60) print("🔍 测试 PyTorch 原始模型") print("="*60) model_path = "/Users/gt/.lmstudio/hub/models/paddleocr-vl" # 加载模型 print("\n📂 加载模型...") model = AutoModel.from_pretrained( model_path, trust_remote_code=True, torch_dtype=torch.float32 ) model.eval() processor = AutoProcessor.from_pretrained( model_path, trust_remote_code=True ) print("✅ 模型加载完成") # 处理图像 print("\n🔄 处理图像...") prompt = "Extract all text from this image." inputs = processor( images=img, text=prompt, return_tensors="pt" ) print(f"pixel_values shape: {inputs['pixel_values'].shape}") print(f"input_ids shape: {inputs['input_ids'].shape}") print(f"image_grid_thw: {inputs['image_grid_thw']}") # 获取中间层输出 print("\n🔍 获取中间层输出...") with torch.no_grad(): # 1. Vision encoder vision_outputs = model.vision_model(inputs['pixel_values']) vision_features = vision_outputs.last_hidden_state print(f"\n1. Vision encoder output: {vision_features.shape}") print(f" 均值: {vision_features.mean():.6f}") print(f" 标准差: {vision_features.std():.6f}") # 2. mlp_AR (vision projection) vision_embeds = model.mlp_AR(vision_features, inputs['image_grid_thw'].tolist()) if isinstance(vision_embeds, list): vision_embeds = vision_embeds[0] print(f"\n2. mlp_AR output: {vision_embeds.shape}") print(f" 均值: {vision_embeds.mean():.6f}") print(f" 标准差: {vision_embeds.std():.6f}") # 3. 完整的 forward outputs = model( input_ids=inputs['input_ids'], pixel_values=inputs['pixel_values'], image_grid_thw=inputs['image_grid_thw'], max_new_tokens=10, do_sample=False ) print(f"\n3. 生成的 token IDs: {outputs[0][:20].tolist()}") return { 'vision_features': vision_features, 'vision_embeds': vision_embeds, 'inputs': inputs } def test_mlx_model(image_path): """测试 MLX 模型""" print("\n" + "="*60) print("🔍 测试 MLX 实现") print("="*60) model_dir = "/Users/gt/.gemini/antigravity/scratch/paddleocr-mlx-conversion" ocr = FinalOptimizedPaddleOCRMLX(model_dir) # 预处理 prompt = "Extract all text from this image." pixel_values, input_ids, image_grid_thw = ocr.preprocess_image(image_path, prompt) print(f"\npixel_values shape: {pixel_values.shape}") print(f"input_ids shape: {input_ids.shape}") print(f"image_grid_thw: {image_grid_thw.tolist()}") # 获取中间层输出 print("\n🔍 获取中间层输出...") # 1. Patch embedding num_patches = pixel_values.shape[0] x = ocr.model.patch_embedding(pixel_values) x = x.reshape(num_patches, ocr.model.vision_hidden_size) x = mx.expand_dims(x, 0) print(f"\n1. Patch embedding: {x.shape}") print(f" 均值: {float(mx.mean(x)):.6f}") print(f" 标准差: {float(mx.std(x)):.6f}") # 2. Position embedding x = x + ocr.model.position_embedding[:num_patches, :] print(f"\n2. After position embedding: {x.shape}") print(f" 均值: {float(mx.mean(x)):.6f}") print(f" 标准差: {float(mx.std(x)):.6f}") # 3. Vision layers (只测试前几层) for i in range(3): x = ocr.model.vision_layers[i](x, None) print(f"\n3.{i+1}. After vision layer {i}: {x.shape}") print(f" 均值: {float(mx.mean(x)):.6f}") print(f" 标准差: {float(mx.std(x)):.6f}") # 完整的 vision encoding vision_embeds = ocr.model.encode_image(pixel_values, image_grid_thw) print(f"\n4. Final vision embeds: {vision_embeds.shape}") print(f" 均值: {float(mx.mean(vision_embeds)):.6f}") print(f" 标准差: {float(mx.std(vision_embeds)):.6f}") # 生成 result = ocr.generate(pixel_values, input_ids, image_grid_thw, max_tokens=10, repetition_penalty=2.0) print(f"\n5. 生成结果: {result}") return { 'vision_embeds': vision_embeds } def compare_outputs(pytorch_results, mlx_results): """对比输出""" print("\n" + "="*60) print("📊 对比 PyTorch vs MLX") print("="*60) # 对比 vision embeds pt_embeds = pytorch_results['vision_embeds'].numpy() mlx_embeds = np.array(mlx_results['vision_embeds']) print(f"\nVision Embeds 对比:") print(f"PyTorch shape: {pt_embeds.shape}") print(f"MLX shape: {mlx_embeds.shape}") if pt_embeds.shape == mlx_embeds.shape: diff = np.abs(pt_embeds - mlx_embeds) print(f"\n差异统计:") print(f" 最大差异: {diff.max():.6f}") print(f" 平均差异: {diff.mean():.6f}") print(f" 相对误差: {(diff.mean() / np.abs(pt_embeds).mean() * 100):.2f}%") # 检查是否接近 if diff.max() < 0.1: print("✅ 输出非常接近！") elif diff.max() < 1.0: print("⚠️ 输出有一些差异") else: print("❌ 输出差异很大！") else: print("❌ 形状不匹配！") def main(): """主函数""" print("\n" + "="*60) print("🔬 PyTorch vs MLX 详细对比") print("="*60) print("目标: 找出准确度低的根本原因") print("="*60) try: # 创建测试图像 print("\n📋 创建测试图像...") image_path, img = create_test_image() print(f"✅ 测试图像: {image_path}") # 测试 PyTorch 模型 pytorch_results = test_pytorch_model(image_path, img) # 测试 MLX 模型 mlx_results = test_mlx_model(image_path) # 对比输出 compare_outputs(pytorch_results, mlx_results) print("\n" + "="*60) print("✅ 对比完成") print("="*60) except Exception as e: print(f"\n❌ 错误: {e}") import traceback traceback.print_exc() if __name__ == "__main__": main()