File size: 2,411 Bytes
fd448dd | 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 76 77 78 79 80 | """Model architecture registry and factory."""
from typing import Dict, Type, Optional, Any
import torch
from taoTrain.core import BaseModel
from taoTrain.config import ModelConfig
# Global registry for model architectures
_ARCHITECTURE_REGISTRY: Dict[str, Type[BaseModel]] = {}
def register_architecture(name: str):
"""Decorator to register a custom model architecture."""
def decorator(cls: Type[BaseModel]):
if name in _ARCHITECTURE_REGISTRY:
raise ValueError(f"Architecture '{name}' is already registered")
_ARCHITECTURE_REGISTRY[name] = cls
return cls
return decorator
def get_registered_architectures() -> Dict[str, Type[BaseModel]]:
"""Get all registered architectures."""
return _ARCHITECTURE_REGISTRY.copy()
def get_model(
config: Any,
device: Optional[torch.device] = None,
) -> BaseModel:
"""
Create a model instance from config.
Args:
config: ModelConfig instance
device: Device to create model on (defaults to CPU)
Returns:
Model instance
"""
if device is None:
device = torch.device('cpu')
# Handle both raw model configs and full training configs used by wrappers.
config_with_arch = config
if not hasattr(config_with_arch, "architecture_type") and hasattr(config, "model"):
config_with_arch = config.model
# Handle both enum and string values
arch_type = config_with_arch.architecture_type
if isinstance(arch_type, str):
arch_name = arch_type
else:
arch_name = arch_type.value
if arch_name not in _ARCHITECTURE_REGISTRY:
raise ValueError(
f"Unknown architecture: {arch_name}. "
f"Available: {list(_ARCHITECTURE_REGISTRY.keys())}"
)
model_class = _ARCHITECTURE_REGISTRY[arch_name]
model = model_class(config).to(device)
return model
def register_builtin_architectures():
"""Register all built-in architectures."""
# Import here to register (avoid circular imports)
from . import transformer # noqa: F401
from . import taonet # noqa: F401
from . import gamma_net # noqa: F401
from . import multimodal_wrapper # noqa: F401
# Auto-register built-in architectures when module is imported
register_builtin_architectures()
|