El framework Once-for-All (OFA) representa un cambio de paradigma en el entrenamiento y despliegue de redes neuronales. Su propuesta principal es entrenar una única red neuronal elástica que abarca una amplia gama de subredes, permitiendo la extracción de configuraciones especializadas optimizadas para restricciones específicas de hardware sin necesidad de reentrenamiento.

Conceptos Clave: Redes Elásticas y Bloques Básicos de Búsqueda

La base de OFA es una red neuronal convolucional con "elasticidad" integrada. Esto significa que dimensiones clave como la profundidad (número de capas), el ancho (número de canales) y el tamaño del kernel pueden variarse de forma controlada dentro de la red principal ya entrenada. Los bloques de construcción fundamentales, como los bloques de convolución móvil (Mobile Blocks), están diseñados para soportar esta búsqueda. El entrenamiento se enfoca en la Accuracy Predictor, una red neuronal que aprende a preedcir el rendimiento de cualquier subred dentro del espacio de búsqueda definido.

Ventajas sobre los Enfoques Convencionales

Los métodos tradicionales de NAS (Neural Architecture Search) requieren múltiples ciclos completos de entrenamiento para evaluar cada arquitectura candidata, lo que consume recursos computacionales y tiempo. OFA mitiga esto mediante un enfoque de "entrenamiento único". Una vez entrenada la red elástica principal, la búsqueda de la arquitectura óptima se convierte en un problema de muestreo y evaluación rápida usando los predictores de eficiencia y precisión, reduciendo drásticamente el costo computacional y la huella de carbono asociada.

Inicio Rápido con el Framework

Para empezar a utilizar OFA, primero debemos configurar el entorno y obtener el código.

1. Configuración del Entorno

# Clonar el repositorio principal (ejemplo usando una réplica)
git clone https://github.com/mit-han-lab/once-for-all.git
cd once-for-all

# Crear un entorno virtual (opcional pero recomendado)
python -m venv venv_ofa
source venv_ofa/bin/activate  # Linux/macOS
# venv_ofa\Scripts\activate  # Windows

# Instalar el paquete OFA y sus dependencias principales
pip install torch torchvision
pip install -e .

2. Entrenamiento de la Red Elástica (SuperNet)

El proceso de entrenamiento principal de la red elástica se puede iniciar con un script dedicado. El siguiente comando inicia el entrenamiento para una tarea de clasificación en ImageNet, utilizando la progresión de elasticidad.

python train.py \
    --task imagenet_classification \
    --model ofa_mbv3_d234_e346_k357_w1.2 \
    --lr 0.025 \
    --epochs 120 \
    --progressive-shrinking

3. Búsqueda y Evaluación de Subredes Especializadas

Con la red entrenada, podemos buscar subredes que cumplan con objetivos de latencia en hardware específico. El siguiente comando realiza una búsqueda evolutiva para encontrar redes optimizadas para diferentes niveles de latencia.

python search_and_evaluate.py \
    --path ./super_networks/ofa_mbv3_d234_e346_k357_w1.2 \
    --latency-setting mobile_cpu \
    --n-generations 50 \
    --population-size 50

Estructura Interna y Módulos del Framework

La implementación de OFA está organizada en módulos funcionales claros:

• Implementación de Redes Neuronales Elásticas: ofa/elastic_nn/network/ define las arquitecturas de red como OFAMobileNetV3.
• Bloques Elásticos: ofa/elastic_nn/modules/ contiene las capas como DynamicConvLayer y DynamicPointLayer que implementan la elasticidad.
• Algoritmos de Búsqueda NAS: ofa/nas/search_algorithm/ incluye algoritmos como la búsqueda evolutiva (EvolutionarySearcher).
• Predictores de Rendimiento: ofa/nas/accuracy_predictor/ y ofa/nas/latency_predictor/ contienen las redes predictoras entrenadas para estimar precisión y eficiencia sin inferencia completa.
• Puntos de Entrada: Los scripts principales como train.py y search_and_evaluate.py sirven como orquestadores del proceso.

Este diseño modular facilita la extensión del framework para nuevos tipos de elasticidad, tareas o predictores de eficiencia para hardware específico como GPUs, CPUs móviles o aceleradores edge.