Liberando una aceleración 2x para DNN: transformando modelos con redes neuronales integrales en solo 1 minuto

Las redes neuronales integrales (INN) son arquitecturas flexibles que una vez entrenadas se pueden transformar a un tamaño arbitrario definido por el usuario sin ningún ajuste. Como las ondas sonoras (música) se pueden muestrear a cualquier frecuencia de muestreo deseada (calidad de sonido), las INN pueden cambiar la forma de los datos y parámetros (calidad DNN) dinámicamente.

El equipo de TheStage.ai presentó su artículo, “Redes neuronales integrales”, en la conferencia IEEE/CVF CVPR 2023. Este trabajo fue reconocido como uno de los 12 artículos 'candidatos al premio' presentados en el evento. Las INN son una nueva clase de redes neuronales que combinan parámetros continuos y operadores integrales para representar capas básicas. Durante la etapa de inferencia, los INN se convierten en una representación DNN básica mediante muestreo discreto a partir de pesos continuos. Los parámetros de dichas redes son continuos a lo largo de las dimensiones del filtro y del canal, lo que conduce a una poda estructurada sin ajuste fino simplemente mediante la rediscretización a lo largo de esas dimensiones.

En este artículo, lo guiaremos a través del proceso de transformación de un modelo EDSR de súper resolución de imagen 4x en un INN y luego le mostraremos cómo lograr una poda estructurada del modelo. Finalmente, volveremos a convertir el INN en un DNN discreto y lo implementaremos en una GPU Nvidia para una inferencia eficiente. Este artículo procederá de la siguiente manera:

1. Una breve introducción a las posadas.
2. Descripción general de la red EDSR para tareas de superresolución.
3. Aplicación framework TorchIntegral para obtener EDSR integral en una sola línea de código.
4. Poda estructural de INN sin ajuste de INN (canalización rápida).
5. Implemente modelos podados en GPU Nvidia.

Para obtener más información y actualizaciones, consulte los siguientes recursos:

Sitio del proyecto TIN

Proyecto POSADA Github

Código de soporte de este artículo.

Poda de INN para DNN sin ajuste fino

Las capas en los INN se reemplazan por operadores integrales, pero la evaluación práctica de los operadores integrales requiere la discretización de las señales de entrada para el uso de métodos de integración numérica. Resulta que las capas en las INN están diseñadas de tal manera que coincidan con las capas DNN clásicas (completamente conectadas, convoluciones) después de la discretización.

Poda de superresolución de imagen EDSR 4x

La superresolución de imágenes es una tarea de visión por computadora bien conocida en la que una imagen debe mejorarse con un operador de degradación conocido o desconocido. Consideramos la forma clásica de superresolución utilizando la reducción de resolución bicúbica como operador de degradación.

Se pueden utilizar muchas arquitecturas para tareas de superresolución de imágenes, incluidas redes neuronales de alta gama basadas en modelos de difusión y transformadores. En este documento, nos centraremos en la arquitectura 4x EDSR. La arquitectura EDSR es adecuada para nuestra demostración, ya que comprende ResNet (que se usa ampliamente en muchos problemas de aprendizaje profundo) y un bloque de muestreo ascendente 4x al final. En las siguientes figuras se puede encontrar una descripción esquemática de EDSR.

Detalles de la poda de la arquitectura EDSR

La poda estructurada, que implica eliminar filtros o canales completos, tiene implicaciones únicas para los bloques residuales, que sirven como componentes básicos de EDSR. Dado que cada estado se actualiza agregando un bloque Conv -> ReLU -> Conv a la entrada, las señales de entrada y salida deben tener la misma cantidad de parámetros. Esto se puede gestionar de manera eficiente en TorchIntegral creando un gráfico de dependencia de poda . La siguiente figura ilustra que la segunda convolución de cada bloque residual forma un solo grupo.

Para podar la segunda convolución en un bloque residual, es necesario podar todas las segundas convoluciones en cada bloque residual. Sin embargo, para una configuración más flexible, podamos los filtros de las primeras convoluciones y, por lo tanto, los canales de las segundas convoluciones, en todos los bloques residuales.

Conversión del modelo EDSR a INN EDSR

Para las conversiones de DNN previamente entrenados, utilizamos nuestro algoritmo especial de permutación de canales de filtro con una interpolación aún más suave. El algoritmo de permutación preserva la calidad del modelo al tiempo que hace que los pesos de los DNN parezcan muestreados a partir de funciones continuas.

 import torch import torchintegral as inn from super_image import EdsrModel # creating 4x EDSR model model = EdsrModel.from_pretrained("eugenesiow/edsr", scale=4).cuda() # Transform model layers to integral. # continous_dims and discrete dims define which dimensions # of parameters tensors should be parametrized continuously # or stay fixed size as in discrete networks. # In our case we make all filter and channel dimensions # to be continuous excluding convolutions of the upsample block. model = inn.IntegralWrapper(init_from_discrete=True)( model, example_input, continuous_dims, discrete_dims ).cuda()

Ajuste de la red de integración: poda estructurada posterior al entrenamiento de DNN

El ajuste de la cuadrícula de integración es la operación que selecciona suavemente (bajo optimización SGD) tensores de parámetros cuyos filtros deben muestrearse para números definidos por el usuario. A diferencia de los métodos de eliminación de filtro/canal, los INN generan filtros que pueden combinar varios filtros discretos debido a la operación de interpolación.

 # Set trainable gird for each integral layer # Each group should have the same grid # During the sum of continuous signals # We need to sample it using the same set of points for group in model.groups: new_size = 224 if 'operator' in group.operations else 128 group.reset_grid(inn.TrainableGrid1D(new_size)) # Prepare model for tuning of integration grid model.grid_tuning() # Start training train(model, train_data, test_data)

Durante nuestros experimentos, descubrimos que 500 imágenes dan el mismo resultado que un conjunto completo de Div2k de 4000 imágenes.

Actuación

El modelo INN resultante se puede convertir fácilmente en un modelo discreto e implementar en cualquier GPU NVIDIA. Proporcionamos cuadros por segundo (FPS) en el RTX A4000 con una resolución de entrada de 64x64. El modelo comprimido alcanza casi el doble de velocidad. Para convertir el modelo INN podado a un modelo discreto, se puede utilizar la siguiente línea de código:

 model = model.transform_to_discrete() # then model can be compiled, for instance # compilation can add an additional 1.4x speedup for inference model = torch.compile(model, backend='cudagraphs') 

Conclusión

En este artículo, presentamos una descripción general del artículo candidato al premio CVPR2023, "Redes neuronales integrales". Se aplicó para la poda posterior al entrenamiento del modelo EDSR 4x, logrando una aceleración de casi 2x con una sola línea de código y un ajuste fino de 1 minuto de la cuadrícula de integración.

En nuestros artículos futuros, presentaremos más aplicaciones de las INN y cubriremos más detalles sobre la implementación eficiente del modelo. Manténganse al tanto:

Sitio del proyecto TIN

Proyecto POSADA Github

Código de soporte de este artículo.

¡Gracias por su atención!