File size: 1,624 Bytes
0695d0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from transformers import AutoModel
import torch
# 你的模型路径
MODEL_PATH = "/home/nashen/deepseek-ocr/DeepSeek-OCR-master/DeepSeek-OCR-vllm/model/"
print(f"⏳ 正在加载模型结构 (不加载权重，速度很快)...")
# 技巧：使用 meta device 或只加载 config，避免消耗显存
try:
# 尝试加载到 CPU，不进行量化，只为了看名字
model = AutoModel.from_pretrained(
MODEL_PATH,
trust_remote_code=True,
device_map="cpu", # 强制 CPU，不占显存
torch_dtype=torch.float16
)
except Exception as e:
print(f"❌ 加载出错: {e}")
exit()
print(f"\n{'='*20} 模型层级结构 {'='*20}")
print("🔍 正在寻找可能的视觉层 (Vision) 和投影层 (Projector)...")
print(f"{'='*60}")
# 遍历所有层
vision_candidates = []
for name, module in model.named_modules():
# 只打印顶级和二级模块，防止刷屏
depth = name.count(".")
if depth <= 2:
# 获取类名，帮助判断它是干嘛的
class_name = module.__class__.__name__
print(f"🔹 层名: {name:<40} | 类型: {class_name}")
# 智能识别可能的视觉层
lower_name = name.lower()
if "vision" in lower_name or "img" in lower_name or "encoder" in lower_name or "proj" in lower_name or "align" in lower_name:
vision_candidates.append(name)
print(f"\n{'='*20} 🎯 建议加入白名单的层 {'='*20}")
print("请把以下名字加入 test_benchmark.py 的 llm_int8_skip_modules 列表中：")
for cand in vision_candidates:
print(f' "{cand}",')
print(f"{'='*60}")
|