====== step1:转换timm的权重- timm 的efficientvit 权重转为paddle格式 ====== #!/usr/bin/env python import os import argparse import paddle import traceback def load_pytorch_weights(model_name): try: import torch import timm except ImportError: raise ImportError("需要安装 torch 和 timm: pip install torch timm") print(f"从 timm 加载 {model_name} 模型...") pt_model = timm.create_model(model_name, pretrained=True) pt_state_dict = pt_model.state_dict() print(f"PyTorch 模型参数数量: {len(pt_state_dict)}") return pt_state_dict def convert_name(pt_name): if 'num_batches_tracked' in pt_name: return None name = pt_name.replace('running_mean', '_mean') name = name.replace('running_var', '_variance') return name def debug_weight_mapping(pt_state_dict, model_name): print(f"\n=== 调试模式: {model_name} ===") aggreg_keys = [k for k in pt_state_dict.keys() if 'aggreg' in k] if aggreg_keys: print(f"发现 {len(aggreg_keys)} 个aggreg相关参数:") for key in aggreg_keys[:10]: print(f" {key}") if len(aggreg_keys) > 10: print(f" ... 还有 {len(aggreg_keys)-10} 个") attention_keys = [k for k in pt_state_dict.keys() if any(x in k for x in ['qkv', 'proj', 'kernel_func'])] if attention_keys: print(f"发现 {len(attention_keys)} 个attention相关参数:") for key in attention_keys[:10]: print(f" {key}") if len(attention_keys) > 10: print(f" ... 还有 {len(attention_keys)-10} 个") print(f"总参数数量: {len(pt_state_dict)}") print("=" * 50) def is_large_model(model_name): return 'efficientvit_l' in model_name.lower() def convert_param(pt_name, pt_param): param = pt_param.cpu().numpy() if 'weight' in pt_name and param.ndim == 2: if 'conv' not in pt_name.lower() and 'norm' not in pt_name.lower(): param = param.T return param def convert_state_dict(pt_state_dict, model_name=None, debug=False): pd_state_dict = {} skipped = [] is_large = is_large_model(model_name) if model_name else False if debug or is_large: debug_weight_mapping(pt_state_dict, model_name) for pt_name, pt_param in pt_state_dict.items(): pd_name = convert_name(pt_name) if pd_name is None: skipped.append(pt_name) continue pd_param = convert_param(pt_name, pt_param) pd_state_dict[pd_name] = pd_param print(f"转换完成: {len(pd_state_dict)} 个参数") print(f"模型类型: {'Large系列' if is_large else 'Base系列'}") if skipped: print(f"跳过的参数 ({len(skipped)}): {skipped[:5]}...") return pd_state_dict def save_weights(pd_state_dict, output_path): os.makedirs(os.path.dirname(output_path), exist_ok=True) paddle.save(pd_state_dict, output_path) print(f"\n权重已保存到: {output_path}") print(f"文件大小: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB") model_keys = [ 'efficientvit_b0.r224_in1k', 'efficientvit_b1.r224_in1k', 'efficientvit_b1.r256_in1k', 'efficientvit_b1.r288_in1k', 'efficientvit_b2.r224_in1k', 'efficientvit_b2.r256_in1k', 'efficientvit_b2.r288_in1k', 'efficientvit_b3.r224_in1k', 'efficientvit_b3.r256_in1k', 'efficientvit_b3.r288_in1k', 'efficientvit_l1.r224_in1k', 'efficientvit_l2.r224_in1k', 'efficientvit_l2.r256_in1k', 'efficientvit_l2.r288_in1k', 'efficientvit_l2.r384_in1k', 'efficientvit_l3.r224_in1k', 'efficientvit_l3.r256_in1k', 'efficientvit_l3.r320_in1k', 'efficientvit_l3.r384_in1k' ] def main(): parser = argparse.ArgumentParser(description='转换 timm EfficientViT 权重到 PaddlePaddle') parser.add_argument('--model', type=str, required=False, help='模型名称,如果不指定则转换所有支持的模型') parser.add_argument('--output', type=str, default='weights/', help='输出目录') parser.add_argument('--debug', action='store_true', help='启用调试模式,显示详细的权重映射信息') args = parser.parse_args() # 如果没有指定模型,则遍历所有支持的模型 if args.model is None or args.model == "": for model_name in model_keys: print(f"\n正在转换模型: {model_name}") try: pt_state_dict = load_pytorch_weights(model_name) pd_state_dict = convert_state_dict(pt_state_dict, model_name, args.debug) output_path = os.path.join(args.output, f'{model_name}.pdparams') save_weights(pd_state_dict, output_path) except Exception as e: print(f"转换 {model_name} 时出错: {str(e)}") print(f"详细堆栈信息:\n{traceback.format_exc()}") continue else: # 单个模型转换 if args.model not in model_keys: print(f"错误: 模型 {args.model} 不在支持的模型列表中") print(f"支持的模型: {model_keys}") return pt_state_dict = load_pytorch_weights(args.model) pd_state_dict = convert_state_dict(pt_state_dict, args.model, args.debug) output_path = os.path.join(args.output, f'{args.model}.pdparams') save_weights(pd_state_dict, output_path) print("\n" + "=" * 60) print("转换完成") print("\n使用方法:") print(f" model = efficientvit_b0()") print(f" state_dict = paddle.load('{output_path}')") print(f" model.set_state_dict(state_dict)") print("=" * 60) if __name__ == '__main__': main()