64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
"""
|
||
模型注册表
|
||
用于解耦模型的调用和实现，支持插件式扩展新模型。
|
||
"""
|
||
|
||
# 训练器注册表
|
||
TRAINER_REGISTRY = {}
|
||
|
||
def register_trainer(name, func):
|
||
"""
|
||
注册一个模型训练器。
|
||
|
||
参数:
|
||
name (str): 模型类型名称 (e.g., 'xgboost')
|
||
func (function): 对应的训练函数
|
||
"""
|
||
if name in TRAINER_REGISTRY:
|
||
print(f"警告: 模型训练器 '{name}' 已被覆盖注册。")
|
||
TRAINER_REGISTRY[name] = func
|
||
print(f"✅ 已注册训练器: {name}")
|
||
|
||
def get_trainer(name):
|
||
"""
|
||
根据模型类型名称获取一个已注册的训练器。
|
||
"""
|
||
if name not in TRAINER_REGISTRY:
|
||
# 在打印可用训练器之前，确保它们已经被加载
|
||
from trainers import discover_trainers
|
||
discover_trainers()
|
||
if name not in TRAINER_REGISTRY:
|
||
raise ValueError(f"未注册的模型训练器: '{name}'. 可用: {list(TRAINER_REGISTRY.keys())}")
|
||
return TRAINER_REGISTRY[name]
|
||
|
||
# --- 预测器注册表 ---
|
||
|
||
# 预测器函数需要一个统一的接口，例如:
|
||
# def predictor_function(model, checkpoint, **kwargs): -> predictions
|
||
|
||
PREDICTOR_REGISTRY = {}
|
||
|
||
def register_predictor(name, func):
|
||
"""
|
||
注册一个模型预测器。
|
||
"""
|
||
if name in PREDICTOR_REGISTRY:
|
||
print(f"警告: 模型预测器 '{name}' 已被覆盖注册。")
|
||
PREDICTOR_REGISTRY[name] = func
|
||
|
||
def get_predictor(name):
|
||
"""
|
||
根据模型类型名称获取一个已注册的预测器。
|
||
如果找不到特定预测器，可以返回一个默认的。
|
||
"""
|
||
return PREDICTOR_REGISTRY.get(name, PREDICTOR_REGISTRY.get('default'))
|
||
|
||
# 默认的PyTorch预测逻辑可以被注册为 'default'
|
||
def register_default_predictors():
|
||
from predictors.model_predictor import default_pytorch_predictor
|
||
register_predictor('default', default_pytorch_predictor)
|
||
# 如果其他PyTorch模型有特殊预测逻辑，也可以在这里注册
|
||
# register_predictor('kan', kan_predictor_func)
|
||
|
||
# 注意：这个函数的调用时机很重要，需要在应用启动时执行一次。
|
||
# 我们可以暂时在 model_predictor.py 导入注册表后调用它。 |