Open-Sourcing BiT: exploración de la formación previa a gran escala para la visión artificial



Un estribillo común para los investigadores de la visión por computadora es que las redes neuronales profundas modernas siempre tienen hambre de más datos etiquetados: las CNN de vanguardia actuales deben ser capacitadas en conjuntos de datos como OpenImages o Places , que consisten en más de 1 millón de imágenes etiquetadas. Sin embargo, para muchas aplicaciones, recopilar esta cantidad de datos etiquetados puede ser prohibitivo para el profesional promedio.

Un enfoque común para mitigar la falta de datos etiquetados para las tareas de visión por computadora es usar modelos que hayan sido previamente entrenados en datos genéricos (por ejemplo, ImageNet ). La idea es que las características visuales aprendidas en los datos genéricos se puedan reutilizar para la tarea de interés. A pesar de que este entrenamiento previo funciona razonablemente bien en la práctica, todavía no alcanza la capacidad de comprender rápidamente nuevos conceptos y comprenderlos en diferentes contextos. En un espíritu similar a cómo BERT y T5 han mostrado avances en el dominio del lenguaje, creemos que la capacitación previa a gran escala puede mejorar el rendimiento de los modelos de visión por computadora.

En " Big Transfer (BiT): Aprendizaje de representación visual general ", diseñamos un enfoque para la capacitación previa efectiva de características generales utilizando conjuntos de datos de imágenes a una escala más allá del estándar de facto ( ILSVRC-2012 ). En particular, destacamos la importancia de elegir adecuadamente las capas de normalización y escalar la capacidad de la arquitectura a medida que aumenta la cantidad de datos previos al entrenamiento. Nuestro enfoque exhibe un rendimiento sin precedentes que se adapta a una amplia gama de nuevas tareas visuales, incluida la configuración de reconocimiento de pocos disparos y el punto de referencia ObjectNet de "mundo real" recientemente introducido. Nos complace compartir los mejores modelos de BiT previamente entrenados en conjuntos de datos públicos, junto con el código en TF2, Jax y PyTorch . Esto permitirá que cualquier persona alcance un rendimiento de vanguardia en su tarea de interés, incluso con solo un puñado de imágenes etiquetadas por clase.

Pre-entrenamiento
Para investigar el efecto de la escala de datos, revisamos las opciones de diseño comunes de la configuración previa al entrenamiento (como las normalizaciones de activaciones y pesos, el ancho / profundidad del modelo y los horarios de entrenamiento) utilizando tres conjuntos de datos: ILSVRC-2012 (1.28M imágenes con 1000 clases), ImageNet-21k (imágenes 14M con ~ 21k clases) y JFT (imágenes 300M con ~ 18k clases). Es importante destacar que con estos conjuntos de datos nos concentramos en el régimen de datos grandes previamente poco explorado.

Primero investigamos la interacción entre el tamaño del conjunto de datos y la capacidad del modelo. Para hacer esto, entrenamos arquitecturas clásicas de ResNet , que funcionan bien, a la vez que son simples y reproducibles. Entrenamos variantes desde el estándar "R50x1" de 50 capas de profundidad hasta el "R152x4" 4x más ancho y de 152 capas de profundidad en cada uno de los conjuntos de datos mencionados anteriormente. Una observación clave es que para beneficiarse de más datos, también es necesario aumentar la capacidad del modelo. Esto se ejemplifica con las flechas rojas en el panel izquierdo de la figura siguiente.
Izquierda : para hacer un uso efectivo de un conjunto de datos más grande para la capacitación previa, es necesario aumentar la capacidad del modelo. Las flechas rojas ejemplifican esto: las arquitecturas pequeñas (punto más pequeño) empeoran cuando se entrena previamente en el ImageNet-21k más grande, mientras que las arquitecturas más grandes (puntos más grandes) mejoran. Derecha : la capacitación previa en un conjunto de datos más grande por sí solo no necesariamente mejora el rendimiento, por ejemplo, cuando se pasa de ILSVRC-2012 al ImageNet-21k relativamente más grande. Sin embargo, al aumentar también el presupuesto computacional y la capacitación durante más tiempo, la mejora del rendimiento es pronunciada.
Una segunda observación, aún más importante, es que la duración del entrenamiento se vuelve crucial. Si uno se entrena previamente en un conjunto de datos más grande sin ajustar el presupuesto computacional y entrenar por más tiempo, es probable que el rendimiento empeore. Sin embargo, al adaptar la programación al nuevo conjunto de datos, las mejoras pueden ser significativas.

Durante nuestra fase de exploración, descubrimos otra modificación crucial para mejorar el rendimiento. Mostramos que reemplazar la normalización por lotes (BN, una capa de uso común que estabiliza el entrenamiento mediante la normalización de las activaciones ) por la normalización de grupo (GN) es beneficioso para el pre-entrenamiento a gran escala. Primero, el estado de BN (media y varianza de las activaciones neuronales) necesita un ajuste entre el preentrenamiento y la transferencia, mientras que el GN no tiene estado, por lo que evita esta dificultad. En segundo lugar, BN utiliza estadísticas a nivel de lote, que se vuelven poco confiables con tamaños de lote pequeños por dispositivo que son inevitables para modelos grandes. Dado que GN no calcula estadísticas a nivel de lote, también evita este problema. Para obtener más detalles técnicos, incluido el uso de una técnica de estandarización de peso para garantizar un comportamiento estable, consulte nuestro documento.
Resumen de nuestra estrategia de pre-entrenamiento: tome una ResNet estándar, aumente la profundidad y el ancho, reemplace BatchNorm (BN) con GroupNorm y Weight Standardization (GNWS), y entrene en un conjunto de datos muy grande y genérico para muchas más iteraciones.
Transferir aprendizaje
Siguiendo los métodos establecidos en el dominio del lenguaje por BERT , ajustamos el modelo BiT pre-entrenado en datos de una variedad de tareas de interés "posteriores", que pueden venir con muy pocos datos etiquetados. Debido a que el modelo pre-entrenado ya viene con una buena comprensión del mundo visual, esta estrategia simple funciona notablemente bien.

El ajuste fino viene con una gran cantidad de hiperparámetros para elegir, como la tasa de aprendizaje, la pérdida de peso, etc. Proponemos una heurística para seleccionar estos hiperparámetros que llamamos "BiT-HyperRule" , que se basa solo en características de conjunto de datos de alto nivel, como la resolución de imagen y el número de ejemplos etiquetados. Aplicamos con éxito la BiT-HyperRule en más de 20 tareas diversas, que van desde imágenes naturales hasta imágenes médicas.
Una vez que el modelo BiT está pre-entrenado, puede ajustarse en cualquier tarea, incluso si solo hay pocos ejemplos etiquetados disponibles.
Al transferir BiT a tareas con muy pocos ejemplos, observamos que a medida que aumentamos simultáneamente la cantidad de datos genéricos utilizados para la capacitación previa y la capacidad de la arquitectura, la capacidad del modelo resultante para adaptarse a datos novedosos mejora drásticamente. Tanto en el CIFAR de 1 disparo como en el de 5 disparos (ver Fig. A continuación), aumentar la capacidad del modelo produce rendimientos limitados cuando se realiza un entrenamiento previo en ILSVRC (curvas verdes). Sin embargo, con el entrenamiento previo a gran escala en JFT, cada aumento en la capacidad del modelo produce rendimientos masivos (curvas marrones), hasta BiT-L, que alcanza el 64% de 1 disparo y el 95% de 5 disparos.
Las curvas representan la precisión media en 5 carreras independientes (puntos de luz) cuando se transfieren a CIFAR-10 con solo 1 o 5 imágenes por clase (10 o 50 imágenes en total). Es evidente que las arquitecturas grandes pre-entrenadas en grandes conjuntos de datos son significativamente más eficientes.
Para verificar que este resultado sea más general, también evaluamos BiT en VTAB-1k , que es un conjunto de 19 tareas diversas con solo 1000 ejemplos etiquetados por tarea. Transferimos el modelo BiT-L a todas estas tareas y logramos un puntaje de 76.3% en general, que es una mejora absoluta de 5.8% sobre el estado del arte anterior .

Mostramos que esta estrategia de pre-entrenamiento a gran escala y transferencia simple es efectiva incluso cuando hay una cantidad moderada de datos disponibles al evaluar BiT-L en varios puntos de referencia estándar de visión por computadora, como Oxford Pets and Flowers , CIFAR , etc. De estos, BiT-L iguala o supera los resultados de vanguardia. Finalmente, usamos BiT como columna vertebral para RetinaNet en la tarea de detección MSCOCO-2017 y confirmamos que incluso para una tarea de salida tan estructurada, el uso de entrenamiento previo a gran escala ayuda considerablemente.
Izquierda: precisión de BiT-L en comparación con el modelo general de vanguardia anterior en varios puntos de referencia estándar de visión por computadora. Derecha: Resultados en la precisión promedio (AP) del uso de BiT como columna vertebral para RetinaNet en MSCOCO-2017.
Es importante enfatizar que en todas las diferentes tareas posteriores que consideramos, no realizamos ajustes de hiperparámetros por tarea y confiamos en BiT-HyperRule . Como mostramos en el documento, se pueden lograr resultados aún mejores ajustando hiperparámetros en datos de validación suficientemente grandes.

Evaluación con ObjectNet
Para evaluar aún más la robustez de BiT en un escenario más desafiante, evaluamos los modelos de BiT que se ajustaron en ILSVRC-2012 en el conjunto de datos ObjectNet recientemente introducido. Este conjunto de datos se parece mucho a los escenarios del mundo real, donde los objetos pueden aparecer en un contexto atípico, punto de vista, rotación, etc. Curiosamente, el beneficio de la escala de datos y arquitectura es aún más pronunciado con el modelo BiT-L que logra una precisión de top 5 sin precedentes de 80.0 %, una mejora absoluta de casi el 25% sobre el estado del arte anterior .
Resultados de BiT en el conjunto de datos de evaluación ObjectNet. Izquierda: precisión de top 5, derecha: precisión de top 1.
Conclusión
Mostramos que, dada la capacitación previa sobre grandes cantidades de datos genéricos, una estrategia de transferencia simple conduce a resultados impresionantes, tanto en grandes conjuntos de datos como en tareas con muy pocos datos, hasta una sola imagen por clase. Nos liberamos el modelo de bit-M, un R152x4 pre-formados en IMAGEnet-21k, junto con colabs para la transferencia de Jax, TensorFlow2 y PyTorch. Además del lanzamiento del código, remitimos al lector al tutorial práctico TensorFlow2 sobre cómo usar los modelos BiT. Esperamos que los profesionales e investigadores encuentren una alternativa útil a los modelos pre-entrenados de ImageNet comúnmente utilizados.

Agradecimientos
Nos gustaría agradecer a Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly y Neil Houlsby, quienes han sido coautores del artículo de BiT y han estado involucrados en todos los aspectos de su desarrollo, así como al equipo de Brain en Zürich. También nos gustaría agradecer a Andrei Giurgiu por su ayuda en la depuración de las tuberías de entrada. Agradecemos a Tom Small por crear las animaciones utilizadas en este blog. Finalmente, remitimos al lector interesado a los enfoques relacionados en esta dirección por parte de nuestros colegas en Google Research, Noisy Student , así como en la muy relevante Exploración de los límites de la capacitación premonitoria supervisada de Facebook Research .