|
|
import torch
|
|
|
import os
|
|
|
from collections import OrderedDict
|
|
|
|
|
|
def prune_and_convert_model():
|
|
|
"""
|
|
|
交互式脚本，用于加载PyTorch检查点，移除trainer状态和text_encoder权重，
|
|
|
并根据用户选择的目标精度（fp32, fp16, bf16）保存。
|
|
|
"""
|
|
|
|
|
|
source_path = input("请输入源模型文件的完整路径（例如 D:\\models\\SDMatte_plus.pth）：\n> ")
|
|
|
|
|
|
if not os.path.exists(source_path):
|
|
|
print(f"\n[错误] 文件未找到: {source_path}")
|
|
|
return
|
|
|
|
|
|
while True:
|
|
|
target_precision = input("\n请选择目标精度 (输入 fp32, fp16, 或 bf16):\n> ").lower()
|
|
|
if target_precision in ['fp32', 'fp16', 'bf16']:
|
|
|
break
|
|
|
else:
|
|
|
print("[错误] 无效输入，请输入 'fp32', 'fp16', 或 'bf16' 中的一个。")
|
|
|
|
|
|
|
|
|
|
|
|
base, ext = os.path.splitext(source_path)
|
|
|
output_path = f"{base}_pruned_{target_precision}{ext}"
|
|
|
|
|
|
print(f"\n--- 开始处理检查点: {os.path.basename(source_path)} ---")
|
|
|
print(f"目标精度: {target_precision.upper()}")
|
|
|
print(f"输出文件将保存为: {os.path.basename(output_path)}")
|
|
|
|
|
|
try:
|
|
|
print("\n正在加载原始权重文件...")
|
|
|
full_checkpoint = torch.load(source_path, map_location="cpu", weights_only=False)
|
|
|
|
|
|
state_dict = None
|
|
|
if isinstance(full_checkpoint, dict):
|
|
|
if 'model' in full_checkpoint:
|
|
|
state_dict = full_checkpoint['model']
|
|
|
print("检测到训练检查点，已从 'model' 键中提取权重。")
|
|
|
elif 'state_dict' in full_checkpoint:
|
|
|
state_dict = full_checkpoint['state_dict']
|
|
|
print("检测到训练检查点，已从 'state_dict' 键中提取权重。")
|
|
|
else:
|
|
|
state_dict = full_checkpoint
|
|
|
print("文件似乎是一个纯权重字典，将直接处理。")
|
|
|
else:
|
|
|
print("[错误] 文件格式无法识别。")
|
|
|
return
|
|
|
|
|
|
pruned_state_dict = OrderedDict()
|
|
|
removed_count = 0
|
|
|
kept_count = 0
|
|
|
|
|
|
print("正在移除 'text_encoder.' 相关权重...")
|
|
|
for key, value in state_dict.items():
|
|
|
if not key.startswith("text_encoder."):
|
|
|
pruned_state_dict[key] = value
|
|
|
kept_count += 1
|
|
|
else:
|
|
|
removed_count += 1
|
|
|
|
|
|
print(f"移除 {removed_count} 个 text_encoder 权重，保留 {kept_count} 个。")
|
|
|
|
|
|
|
|
|
if target_precision != 'fp32':
|
|
|
print(f"开始将保留的权重转换为 {target_precision} ...")
|
|
|
target_dtype = torch.float16 if target_precision == 'fp16' else torch.bfloat16
|
|
|
|
|
|
for key, value in pruned_state_dict.items():
|
|
|
if isinstance(value, torch.Tensor) and value.is_floating_point():
|
|
|
pruned_state_dict[key] = value.to(target_dtype)
|
|
|
else:
|
|
|
print("保留原始 FP32 精度。")
|
|
|
|
|
|
print(f"正在保存处理后的模型到: {output_path}")
|
|
|
torch.save(pruned_state_dict, output_path)
|
|
|
|
|
|
source_size = os.path.getsize(source_path) / (1024**2)
|
|
|
pruned_size = os.path.getsize(output_path) / (1024**2)
|
|
|
|
|
|
print("\n--- 成功 ---")
|
|
|
print(f"原始文件大小: {source_size:.2f} MB")
|
|
|
print(f"裁剪后文件大小 ({target_precision.upper()}): {pruned_size:.2f} MB")
|
|
|
print(f"总共节省空间: {source_size - pruned_size:.2f} MB")
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"\n[错误] 处理过程中发生错误: {e}")
|
|
|
import traceback
|
|
|
traceback.print_exc()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
prune_and_convert_model() |
