对funasr微调的模型进行压缩

funasr语音识别模型原本是840M

模型微调后除了模型的权重文件还包括以下内容

训练配置（config）
优化器状态（optimizer）
学习率调度器状态（scheduler）
训练进度（epoch、step 等）
随机数种子状态

因此训练后模型变为2.4个G

通过以下代码将训练后的模型进行修改和压缩

import torch
import os

def convert_to_half(obj):
    """递归地将张量转换为半精度"""
    if torch.is_tensor(obj) and torch.is_floating_point(obj):
        return obj.half()
    elif isinstance(obj, dict):
        return {k: convert_to_half(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_half(v) for v in obj]
    else:
        return obj

# 加载模型
ckpt = torch.load(
"../examples/industrial_data_pretraining/paraformer_streaming/outputs/model.pt.best",
    map_location="cpu"
)

# FunASR 一般把权重放在 'model' 或 'state_dict' 字段里
if "model" in ckpt:
    state_dict = ckpt["model"]
elif "state_dict" in ckpt:
    state_dict = ckpt["state_dict"]
else:
    raise KeyError("未找到模型权重字段，请检查 ckpt 结构")


type = input('''
请输入压缩类型
1、只删除训练数据文件
2、模型转为半精度float16
''')

if int(type) == 1:
    modelname = 'model_weights.pt'
    # 保存成纯权重文件
    torch.save(state_dict, modelname)
else:
    modelname = 'model_fp16.pt'
    # 转成半精度
    model_fp16 = convert_to_half(state_dict)
    # 保存
    torch.save(model_fp16, modelname)

file_path = os.path.abspath(__file__)
dir_path = os.path.dirname(file_path)
full_path = os.path.join(dir_path, modelname)

print('已保存到：',full_path)

以上代码分两部分，去掉训练数据只保留模型，将模型转为半精度，

按提示输入对应数字进行压缩，会保存到对应目录下。