4.3. Modelos y registry (src/models/)

Responsabilidad: dado un nombre-string, devolver un nn.Module listo con su cabeza adaptada a las 9 clases. Tres piezas: registry.py (el patrón), baselines/*.py (las 8 arquitecturas), input_sizes.py (excepciones de resolución).

registry.py — el patrón Registry

Propósito y por qué existe

Desacoplar quién pide un modelo de cómo se construye. El código de entrenamiento solo conoce strings ("efficientnet_b0"); no importa ni instancia arquitecturas concretas. Añadir un modelo nuevo no requiere tocar el código de entrenamiento: basta con registrarlo.

Clase ModelRegistry

Es esencialmente un diccionario nombre → ModelEntry(name, factory, default_kwargs) con métodos de conveniencia.

def build(self, name: str, num_classes: int, **override_kwargs) -> nn.Module:
    if name not in self._entries:
        raise KeyError(f"Modelo '{name}' no registrado. Disponibles: {available}")
    entry = self._entries[name]
    kwargs = {**entry.default_kwargs, **override_kwargs, "num_classes": num_classes}
    return entry.factory(**kwargs)

Bloque clave — el merge de kwargs: el orden {**default, **override, "num_classes": ...} significa que los override_kwargs del llamador pisan los defaults del registro, y num_classes siempre gana (no puede ser sobrescrito accidentalmente). Es la fusión de configuración por capas, con la más específica al final.

Instancia global: MODEL_REGISTRY = ModelRegistry() — un singleton de módulo. Los archivos de baselines/ lo importan y se auto-registran al importarse.

Caso borde bien manejado: pedir un modelo inexistente lanza KeyError con la lista de disponibles, no un error críptico. Mismo patrón en get_entry.

El mecanismo de auto-registro (importante entenderlo)

En src/models/__init__.py:

import src.models.baselines.efficientnet  # noqa: F401 - registers models
import src.models.baselines.fastvit       # noqa: F401
import src.models.baselines.ghostnet      # noqa: F401
import src.models.baselines.mobilenet     # noqa: F401
import src.models.baselines.shufflenet    # noqa: F401

Intención: el simple hecho de importar cada módulo ejecuta sus llamadas MODEL_REGISTRY.register(...) de nivel superior, poblando el registry. El comentario # noqa: F401 silencia el aviso del linter de "import sin usar" — el import sí se usa, por su efecto secundario. Por eso los scripts de explicabilidad repiten estos imports: necesitan que el registry esté poblado antes de construir modelos.

baselines/*.py — las 8 arquitecturas registradas

Cada archivo define funciones-fábrica y las registra. Hay dos estilos de construcción según la fuente de los pesos:

Estilo torchvision (EfficientNet-B0, ShuffleNetV2)

def build_efficientnet_b0(num_classes, pretrained=True, **kwargs) -> nn.Module:
    weights = tv_models.EfficientNet_B0_Weights.DEFAULT if pretrained else None
    model = tv_models.efficientnet_b0(weights=weights, **kwargs)
    in_features = model.classifier[1].in_features
    model.classifier[1] = nn.Linear(in_features, num_classes)  # ⭐ reemplaza la cabeza
    return model

El bloque crítico es el reemplazo de la cabeza: el modelo preentrenado clasifica 1000 clases de ImageNet; se lee cuántas features entran a la última capa lineal (in_features) y se sustituye por una nueva Linear de num_classes salidas (9). Todo lo demás (el backbone convolucional) conserva sus pesos preentrenados. Esto es transfer learning en la práctica.

ShuffleNetV2 hace lo mismo pero su cabeza se llama model.fc en vez de model.classifier[1].

Estilo timm (EfficientNet-Lite0/B4, MobileNetV3, GhostNetV2, FastViT)

def build_efficientnet_lite0(num_classes, pretrained=True, **kwargs) -> nn.Module:
    import timm
    return timm.create_model("efficientnet_lite0", pretrained=pretrained,
                             num_classes=num_classes, **kwargs)

Por qué timm: timm trae cientos de arquitecturas modernas que torchvision no incluye (EfficientNet-Lite, GhostNetV2, FastViT). Su create_model(..., num_classes=N) ya reemplaza la cabeza internamente, así que no hay que hacerlo a mano. El import timm es diferido (dentro de la función) para no forzar la dependencia si solo usas modelos de torchvision.

Catálogo completo

Nombre registrado	Fuente	Tipo	Notas
efficientnet_b0	torchvision	CNN	Líder en macro-F1 (0.9117)
efficientnet_lite0	timm	CNN	Variante "lite" pensada para móvil/edge
efficientnet_b4	timm	CNN	Más grande; entrada nativa 380
mobilenet_v3_large	timm	CNN	Clásico de edge
mobilenet_v3_small	timm	CNN	Aún más ligero
shufflenet_v2_x1_0	torchvision	CNN	Muy eficiente; macro-F1 0.9083
ghostnetv2_100	timm	CNN	Convoluciones "fantasma" baratas
fastvit_t8	timm	Híbrido CNN+Transformer	Entrada nativa 256

Los 3 baselines por defecto (los que corre make train-baselines) son efficientnet_b0, shufflenet_v2_x1_0, efficientnet_lite0: un equilibrio entre precisión y ligereza, todos candidatos realistas para el despliegue móvil.

input_sizes.py — resolución nativa por modelo

Propósito

Algunas arquitecturas fueron diseñadas/preentrenadas a una resolución distinta de 224. Forzarlas a 224 desperdicia capacidad (B4) o penaliza los pesos preentrenados (FastViT).

MODEL_INPUT_SIZES: dict[str, tuple[int, int]] = {
    "efficientnet_b4": (380, 380),  # diseñado para 380x380
    "fastvit_t8": (256, 256),       # preentrenado a 256x256
}

def resolve_input_size(name, fallback):
    return MODEL_INPUT_SIZES.get(name, fallback)

Diseño minimalista: solo se listan las excepciones. Cualquier modelo ausente usa el fallback (el target_size del YAML, 224). Cuando el entrenamiento resuelve una resolución mayor, train_baselines.py reduce automáticamente el batch inversamente al área para no exceder la memoria de la GPU (ver entrenamiento).

Cómo añadir un modelo nuevo (checklist)

Si en el futuro registras una arquitectura, tienes que tocar dos lugares (fácil de olvidar el segundo):

1. src/models/baselines/<archivo>.py: define build_xxx y MODEL_REGISTRY.register("xxx", ...). Si no está en torchvision, usa timm con import diferido.
2. src/explainability/gradcam.py → GRADCAM_TARGET_LAYERS: añade la entrada "xxx": "<ruta.a.la.última.capa.espacial>". Si lo olvidas, Grad-CAM se omite con un warning y el reporte cae a 3 paneles (LIME solo) en vez de 4. Ver explicabilidad.
3. Si la arquitectura usa una resolución nativa distinta de 224, añádela a MODEL_INPUT_SIZES.
4. Si es de torchvision, asegúrate de importarla en src/models/__init__.py (los de baselines/ ya se importan por archivo).

Siguiente

Ya tienes datos y modelos. Ahora el pegamento: el Entrenamiento.