File size: 2,558 Bytes
eb10636 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import os
from rkllm.api import RKLLM
def convert_model(model_path, output_name, do_quantization=False):
"""转换单个模型"""
llm = RKLLM()
print(f"正在加载模型: {model_path}")
ret = llm.load_huggingface(model=model_path, model_lora=None, device='cpu')
if ret != 0:
print(f'加载模型失败: {model_path}')
return ret
print(f"正在构建模型: {output_name} (量化: {do_quantization})")
qparams = None
ret = llm.build(do_quantization=do_quantization, optimization_level=1, quantized_dtype='w8a8',
quantized_algorithm='normal', target_platform='rk3588', num_npu_core=3, extra_qparams=qparams)
if ret != 0:
print(f'构建模型失败: {output_name}')
return ret
# 导出rkllm模型
print(f"正在导出模型: {output_name}")
ret = llm.export_rkllm(output_name)
if ret != 0:
print(f'导出模型失败: {output_name}')
return ret
print(f"成功转换: {output_name}")
return 0
def main():
"""主函数：遍历所有子文件夹并转换模型"""
current_dir = '.'
# 获取所有子文件夹
subdirs = [d for d in os.listdir(current_dir)
if os.path.isdir(os.path.join(current_dir, d)) and not d.startswith('.')]
print(f"找到 {len(subdirs)} 个模型文件夹: {subdirs}")
for subdir in subdirs:
model_path = os.path.join(current_dir, subdir)
# 生成输出文件名
base_name = subdir.replace('/', '_').replace('\\', '_')
quantized_output = f"{base_name}_w8a8.rkllm"
unquantized_output = f"{base_name}_f16.rkllm"
print(f"\n{'='*50}")
print(f"处理模型文件夹: {subdir}")
print(f"{'='*50}")
# 转换非量化版本
print(f"\n--- 转换非量化版本 ---")
ret = convert_model(model_path, unquantized_output, do_quantization=False)
if ret != 0:
print(f"非量化版本转换失败: {subdir}")
continue
# 转换量化版本
print(f"\n--- 转换量化版本 ---")
ret = convert_model(model_path, quantized_output, do_quantization=True)
if ret != 0:
print(f"量化版本转换失败: {subdir}")
continue
print(f"\n✓ {subdir} 模型转换完成!")
print(f" - 非量化版本: {unquantized_output}")
print(f" - 量化版本: {quantized_output}")
if __name__ == "__main__":
main()
|