CodeJeNN: A simple C++ neural network generator for physics applications

Resumen Técnico: CodeJeNN

Planteamiento del Problema
El aprendizaje automático ha demostrado un potencial significativo para acelerar los métodos numéricos en la física, particularmente en la dinámica de fluidos computacional (CFD) para tareas como el modelado de turbulencia y modelos constitutivos sustitutos. Sin embargo, existe un cuello de botella crítico en la integración de las librerías de aprendizaje automático basadas en Python (por ejemplo, TensorFlow, PyTorch, Keras) en los solvers de C++ o Fortran compilados de alto rendimiento. Esta disparidad introduce una sobrecarga computacional y complica el despliegue de redes neuronales en marcos de trabajo de física establecidos. Además, depender de las APIs de C++ proporcionadas por los frameworks suele limitar las capacidades de optimización, ya que fuerza las operaciones tensoriales hacia librerías externas que pueden entrar en conflicto con las estrategias existentes de paralelización y vectorización. Si bien soluciones previas como RoseNNa existen para Fortran y modelos ONNX, sigue habiendo una necesidad de una herramienta que genere código C++ autónomo y libre de dependencias, específicamente para modelos Keras, para facilitar una integración fluida sin dependencias de tiempo de ejecución externas.

Metodología
Los autores presentan CodeJeNN, un generador de redes neuronales en C++ que cierra la brecha entre los modelos Keras entrenados y los solvers de física compilados. El flujo de trabajo opera de la siguiente manera:

1. Entrenamiento del Modelo: Los usuarios entrenan una red neuronal (MLP o CNN) utilizando Keras y guardan la arquitectura y los parámetros en un archivo .keras o .h5.
2. Generación de Código: CodeJeNN carga el modelo, extrae los parámetros en un diccionario de Python y reconstruye la lógica del modelo.
3. Síntesis de C++: La herramienta genera un archivo de cabecera C++ autónomo. Los principios de diseño clave incluyen:
   • Inlining (Incorporación): Inlining agresivo de las funciones de activación y propagación para minimizar la sobrecarga de llamadas a funciones.
   • Reducción de Dependencias: El código generado utiliza únicamente librerías estándar de C++ (por ejemplo, <iostream>, <array>, <cmath>), eliminando la necesidad de librerías de álgebra lineal externas como LAPACK o Eigen.
   • Optimización en Tiempo de Compilación: El uso de las palabras clave constexpr, static e inline asegura que los tamaños de los arreglos y las operaciones se resuelvan en el tiempo de compilación, optimizando la utilización de la caché.
   • Extensibilidad: La estructura del código fuente está diseñada para permitir que los Modelos de Lenguaje de Gran Escala (LLM) generen la lógica para arquitecturas no soportadas con un consumo mínimo de tokens, facilitando la extensión no invasiva del framework.

La herramienta soporta una amplia gama de capas de Keras (Dense, Convolutional, Pooling, Normalization, etc.) y funciones de activación (ReLU, GELU, Swish, etc.), incluyendo activaciones personalizadas definidas por el usuario, siempre que estén serializadas y definidas en un script de Python.

Contribuciones Clave

• Generación de Código Autónomo: CodeJeNN produce un único archivo de cabecera C++ que contiene todos los pesos, sesgos y la lógica necesarios, permitiendo su inclusión directa en solvers existentes sin dependencias de tiempo de ejecución externas.
• Optimización del Rendimiento: Al evitar las llamadas a APIs externas y aprovechar el inlining agresivo, el código generado está diseñado para ejecutarse eficientemente en el hardware local.
• Capacidades de Depuración: La herramienta incluye una API de depuración que compara el modelo extraído contra el model.summary() de Keras y, cuando está habilitada, imprime las salidas de las capas intermedias para ayudar a localizar errores de lógica o algebraicos durante la inferencia.
• Extensibilidad vía LLM: El código base está estructurado para minimizar el costo de tokens y el esfuerzo requerido para añadir soporte para nuevas arquitecturas utilizando LLMs.

Resultos
El artículo valida CodeJeNN a través de dos benchmarks principales:

1. Benchmarks de Velocidad de Inferencia: Las pruebas se realizaron en un Apple M1 Pro comparando CodeJeNN contra Keras ejecutándose en backends de TensorFlow y PyTorch (tanto en modos eager como JIT-compiled).
   • Para un Perceptrón Multicapa (MLP), CodeJeNN logró una aceleración de 25x sobre el TensorFlow estándar y una aceleración de 38x sobre el PyTorch estándar.
   • Incluso contra modelos Keras compilados con JIT, CodeJeNN se mantuvo aproximadamente 2.5x más rápido (TensorFlow) y 9x más rápido (PyTorch).
   • Para Redes Neuronales Convolucionales (CNN), las aceleraciones fueron sustanciales (18x–26x sobre Keras estándar), aunque el rendimiento decayó para modelos más grandes al transicionar de estar limitados por cómputo a estar limitados por memoria debido a las limitaciones de caché.
2. Aplicación en CFD (Modelado de Viscosidad): Una red neuronal se integró en un solver basado en OpenFOAM (DetonationFOAM) para reemplazar la regla de mezcla de Wilke para calcular la viscosidad en una capa de mezcla de hidrógeno-aire.
   • La red, entrenada en 10,000 estados, se incrustó directamente en los encabezados C++ del solver.
   • En las simulaciones, el sustituto de la red neuronal introdujo un error máximo de 0.25% en el campo de velocidad en comparación con el modelo nativo de Wilke, preservando la evolución general del flujo y el desarrollo de la inestabilidad.

Significancia y Reivindicaciones
Los autores afirman que CodeJeNN reduce significativamente la barrera para que los investigadores de física integren redes neuronales en solvers de alto rendimiento. Al generar código fuente bruto e inlined, elimina la necesidad de una infraestructura compleja o la dependencia de librerías de tensores externas, lo cual es crucial para códigos que utilizan todos los recursos de hardware disponibles. La herramienta aborda la ineficiencia de usar LLMs para generar repetidamente código de infraestructura al proporcionar un pipeline dedicado y repetible. El artículo posiciona a CodeJeNN como una solución práctica para la computación científica donde los modelos deben permanecer compactos, eficientes y libres de dependencias externas para satisfacer restricciones de tiempo real y de hardware. El trabajo futuro incluye la integración de soporte para Redes Neuronales de Grafos (GNN) vía Spektral y la exploración de exportaciones de librerías de bajo nivel (por ejemplo, Eigen) para una mayor aceleración.