Dentro de cuda-convnet: las capas de AlexNet en dos GPU y la era de ImageNet

El hito de AlexNet en 2012 se apoyó en ingeniería de GPU de bajo nivel (cuda-convnet), una arquitectura dividida entre dos GPU, y desencadenó directamente una rápida carrera hacia modelos de visión por computadora más profundos y eficientes.

1. Cómo funcionaba cuda-convnet por dentro

Antes de que NVIDIA creara cuDNN, Alex Krizhevsky escribió cuda-convnet directamente en CUDA C++ para extraer la máxima velocidad de las primeras GPU con arquitectura Fermi (la GTX 580). Las bibliotecas modernas suelen desplegar las convoluciones en enormes multiplicaciones de matrices (im2col), lo que consume cantidades ingentes de memoria. Como cada GTX 580 solo tenía 3 GB de VRAM, Krizhevsky diseñó rutinas directas para GPU que ahorraban memoria.

• Disposición de memoria con prioridad al lote (CHWN). Los entornos de trabajo modernos usan por defecto NCHW (lote, canales, altura, anchura), que sitúa todos los píxeles de una misma imagen juntos en memoria. Krizhevsky almacenaba los tensores en orden CHWN, colocando la dimensión del lote (N) al final. Cuando 32 hilos de GPU de un warp procesaban 32 imágenes distintas en la misma posición de píxel, sus lecturas de memoria eran completamente contiguas (coalescentes), lo que maximizaba el ancho de banda de la VRAM.

• Procesamiento por bloques en memoria compartida. Leer desde la VRAM global requiere cientos de ciclos de reloj, mientras que acceder a la memoria compartida del chip (solo 48 KB por Streaming Multiprocessor en la GTX 580) requiere apenas unos pocos ciclos. Los bloques de hilos cargaban pequeñas regiones 2D de píxeles de imagen y pesos de filtros en la memoria compartida, calculaban decenas de productos escalares y después escribían los resultados finales de vuelta en la VRAM.

• Reutilización de registros a nivel de warp. Krizhevsky asignaba a cada bloque de hilos el cálculo de un bloque de píxeles de salida para varios filtros simultáneamente. Al mantener los pesos de los filtros en los rápidos registros de la GPU, una misma región de imagen cargada desde la memoria compartida se reutilizaba en muchos filtros a la vez, lo que reducía drásticamente el tráfico de memoria.

2. Capas específicas de AlexNet y división entre dos GPU

AlexNet contenía 8 capas con parámetros aprendidos (5 convolucionales y 3 totalmente conectadas), que sumaban aproximadamente 60 millones de parámetros. Como una sola GTX 580 no disponía de los más de 3 GB de VRAM necesarios para las activaciones, los pesos y los gradientes, Krizhevsky dividió la red verticalmente en dos flujos paralelos.

• Especialización natural de características. Curiosamente, como las capas 1 y 2 nunca se comunicaban entre GPU, las dos tarjetas desarrollaban sistemáticamente funciones especializadas. Una GPU aprendía detectores de bordes y texturas independientes del color (filtros similares a los de Gabor en escala de grises), mientras que la otra se especializaba en manchas y tonalidades de colores específicos.

• Reducción de los cuellos de botella de PCIe. Transferir datos por el bus PCIe de la placa base era lento. Al restringir la comunicación entre GPU estrictamente a Conv 3 y a las capas totalmente conectadas, Krizhevsky mantenía las GPU calculando a máxima velocidad sin detenerse por las transferencias de datos.

3. Qué ocurrió después en la era de ImageNet

La victoria de AlexNet desencadenó una carrera anual en la competición de ImageNet, que llevó las redes de 8 capas a más de 150 en solo tres años.

• 2013 - ZFNet (visualización y ajuste). Matthew Zeiler y Rob Fergus introdujeron técnicas de deconvolución para visualizar lo que AlexNet veía realmente dentro de sus capas ocultas. Al reducir el tamaño del filtro de la primera capa de AlexNet de 11×11 a 7×7 y disminuir el paso, ZFNet rebajó la tasa de error top-5 de ImageNet del 15,3% al 14,8%.

• 2014 - VGGNet (profundidad mediante filtros uniformes de 3×3). Investigadores de Oxford demostraron que apilar muchas convoluciones pequeñas de 3×3 (hasta 16 o 19 capas) funcionaba mejor que usar filtros grandes individuales (como los de 11×11 o 5×5). Dos capas de 3×3 apiladas cubren la misma área de visión de 5×5 con menos parámetros y activaciones ReLU no lineales adicionales, lo que redujo las tasas de error al 7,3%.

• 2014 - GoogLeNet / Inception (eficiencia multiescala). Google presentó el módulo Inception, que ejecutaba convoluciones de 1×1, 3×3 y 5×5 en paralelo dentro de la misma capa y combinaba sus salidas. Al usar convoluciones de 1×1 de bajo coste para comprimir la profundidad de canales antes de los cálculos costosos, GoogLeNet alcanzó 22 capas y una tasa de error del 6,7% con 12 veces menos parámetros que AlexNet.

• 2015 - ResNet (conexiones de salto y 152 capas). El simple hecho de apilar más capas acabó degradando la precisión del entrenamiento porque los gradientes se desvanecían durante la retropropagación. Microsoft Research lo resolvió con conexiones residuales (conexiones de salto), que permitían a los datos omitir capas mediante atajos directos de identidad. Esto hizo posibles las redes de 152 capas, que lograron una tasa de error top-5 del 3,57%, superando por primera vez el rendimiento humano en la prueba de referencia (~5%).