别再手动写if-else了！用Python装饰器@register_model实现模型工厂，5行代码搞定动态加载
用Python装饰器构建模型工厂：告别if-else的优雅实践
在深度学习项目中，我们经常需要管理多个模型架构。传统做法是通过冗长的if-else语句或配置文件来实例化不同模型，这不仅让代码变得臃肿，还增加了维护成本。想象一下，每次新增一个模型都要修改工厂函数，这种硬编码方式既容易出错又难以扩展。有没有更优雅的解决方案？
Python装饰器提供了一种声明式的编程范式，配合设计模式中的注册表(Registry)机制，我们可以构建一个灵活的模型工厂系统。这种方案特别适合以下场景：
- 多模型对比实验
- A/B测试框架
- 可插拔的模型组件系统
- 需要动态加载模型的应用程序
1. 装饰器与注册表模式的核心原理
装饰器本质上是一个高阶函数，它接受一个函数或类作为输入，并返回修改后的版本。在模型注册场景中，我们可以利用装饰器自动将模型类添加到全局注册表中。
model_registry = {} def register_model(name=None): def decorator(cls): key = name if name else cls.__name__ model_registry[key] = cls return cls return decorator这个register_model装饰器的工作原理是：
- 接受一个可选的
name参数，允许自定义注册名称 - 返回真正的装饰器函数
decorator - 装饰器将目标类与指定名称（或类名）关联存入
model_registry - 最终返回原始类，保持其原有行为不变
与传统的工厂模式相比，这种方案有三大优势：
| 特性 | 传统工厂模式 | 装饰器注册模式 |
|---|---|---|
| 新增模型复杂度 | 高（需修改工厂函数） | 低（只需添加装饰器） |
| 代码耦合度 | 高 | 低 |
| 可读性 | 中等 | 高 |
2. 实现一个生产级模型工厂系统
基础版本虽然简单，但在实际项目中还需要考虑更多细节。下面是一个增强版的实现：
from typing import Dict, Type, Any from torch import nn class ModelRegistry: _registry: Dict[str, Type[nn.Module]] = {} @classmethod def register(cls, name: str = None): def decorator(model_class: Type[nn.Module]): key = name or model_class.__name__ if key in cls._registry: raise ValueError(f"Model {key} already registered") cls._registry[key] = model_class return model_class return decorator @classmethod def create(cls, name: str, *args, **kwargs) -> nn.Module: if name not in cls._registry: available = ", ".join(cls._registry.keys()) raise ValueError(f"Unknown model: {name}. Available: {available}") return cls._registry[name](*args, **kwargs)这个改进版本增加了以下特性：
- 使用类方法而非全局变量，避免命名空间污染
- 类型注解提升代码可读性
- 重复注册检查防止意外覆盖
- 更友好的错误提示，列出可用模型
实际使用时，模型定义变得非常直观：
@ModelRegistry.register("alexnet") class AlexNet(nn.Module): def __init__(self, num_classes: int = 1000): super().__init__() # 网络结构定义... @ModelRegistry.register() class ResNet(nn.Module): # 默认使用类名作为键 def __init__(self, block_type: str, layers: list): super().__init__() # 网络结构定义...3. 高级应用场景与技巧
3.1 动态模型选择与配置
在实验框架中，我们经常需要根据配置文件动态选择模型：
import yaml with open("config.yaml") as f: config = yaml.safe_load(f) model = ModelRegistry.create( config["model"]["name"], **config["model"]["params"] )对应的YAML配置示例：
model: name: "resnet50" params: block_type: "bottleneck" layers: [3, 4, 6, 3] num_classes: 10003.2 模型变体自动注册
通过类继承和装饰器的组合，可以轻松创建并注册模型变体：
@ModelRegistry.register("resnet18") class ResNet18(ResNet): def __init__(self, num_classes=1000): super().__init__("basic", [2, 2, 2, 2], num_classes) @ModelRegistry.register("resnet34") class ResNet34(ResNet): def __init__(self, num_classes=1000): super().__init__("basic", [3, 4, 6, 3], num_classes)3.3 多后端支持
注册表模式可以扩展为支持不同框架的模型：
@ModelRegistry.register("vit_torch") class ViTTorch(nn.Module): # PyTorch实现... @ModelRegistry.register("vit_tf") class ViTTF: # TensorFlow实现... def __call__(self, inputs): # 兼容PyTorch的调用方式 return self.forward(inputs)4. 性能优化与最佳实践
虽然装饰器方案带来了极大便利，但在大型项目中仍需注意以下要点：
线程安全考虑：
- 注册过程通常发生在模块导入时（主线程）
- 如果动态注册，需要加锁保护注册表
from threading import Lock class ThreadSafeModelRegistry(ModelRegistry): _lock = Lock() @classmethod def register(cls, name: str = None): def decorator(model_class: Type[nn.Module]): with cls._lock: return super().register(name)(model_class) return decorator内存优化技巧：
- 延迟加载：只存储类引用而非实例
- 按需导入：在注册装饰器中动态导入大型模块
- 缓存实例：对常用模型添加实例缓存层
class CachedModelRegistry(ModelRegistry): _instance_cache: Dict[str, nn.Module] = {} @classmethod def create(cls, name: str, *args, **kwargs) -> nn.Module: cache_key = f"{name}_{hash(frozenset(kwargs.items()))}" if cache_key not in cls._instance_cache: cls._instance_cache[cache_key] = super().create(name, *args, **kwargs) return cls._instance_cache[cache_key]项目结构建议：
models/ ├── __init__.py # 导出注册表和常用模型 ├── registry.py # 注册表实现 ├── classification/ │ ├── resnet.py │ ├── vit.py │ └── ... └── segmentation/ ├── unet.py └── deeplab.py在__init__.py中集中导入所有模型，确保注册完成：
from .registry import ModelRegistry from .classification import * from .segmentation import * __all__ = ['ModelRegistry']