U-Net es una arquitectura de red neuronal convolucional desarrollada para la segmentación de imágenes, particularmente en contextos biomédicos. Fue introducida en 2015 por Olaf Ronneberger, Philipp Fischer y Thomas Brox en la Universidad de Friburgo, basándose en redes totalmente convolucionales anteriores. La arquitectura se distingue por su diseño simétrico de codificador-decodificador con conexiones de salto, lo que le permite producir mapas de segmentación precisos a partir de relativamente pocas imágenes de entrenamiento. La segmentación de una imagen de 512 × 512 tarda menos de un segundo en una GPU moderna (2015) utilizando la arquitectura U-Net.
La arquitectura U-Net proviene de la llamada "red totalmente convolucional" (FCN), que fue desarrollada por Evan Shelhamer, Jonathan Long y Trevor Darrell en 2014. La idea principal es complementar una red de contracción habitual con capas sucesivas, donde las operaciones de agrupación se reemplazan por operadores de aumento de resolución. Por lo tanto, estas capas aumentan la resolución de la salida. Una capa convolucional sucesiva puede entonces aprender a ensamblar una salida precisa basándose en esta información.
Arquitectura de la Red
La red consta de un camino de contracción y un camino de expansión, lo que le da la arquitectura en forma de U. El camino de contracción es una red convolucional típica que consiste en la aplicación repetida de convoluciones, cada una seguida de una unidad lineal rectificada (ReLU) y una operación de agrupación máxima. Durante la contracción, la información espacial se reduce mientras que la información de características aumenta. El camino de expansión combina la información de características y espacial a través de una secuencia de deconvoluciones y concatenaciones con características de alta resolución del camino de contracción.
Una modificación importante en U-Net es que hay un gran número de canales de características en la parte de aumento de resolución, lo que permite a la red propagar información de contexto a capas de mayor resolución. Como consecuencia, el camino de expansión es más o menos simétrico a la parte de contracción, y produce una arquitectura en forma de U. La red solo utiliza la parte válida de cada convolución sin ninguna capa totalmente conectada. Por lo tanto, la red puede entrenarse y ejecutarse con una mezcla de resoluciones de entrada.
Para predecir los píxeles en la región fronteriza de la imagen, el contexto faltante más allá del límite de la imagen se extrapola reflejando la imagen de entrada en los límites de la imagen. El artículo original de U-Net sugirió además una estrategia de mosaico, donde las imágenes grandes se cortan en mosaicos superpuestos que se procesan de forma independiente. Esto permite procesar imágenes de alta resolución que de otro modo excederían la memoria GPU disponible. Recientemente, también ha habido interés en modelos U-Net basados en campos receptivos para la segmentación de imágenes médicas.
Entrenamiento y Función de Pérdida
U-Net se entrena de extremo a extremo utilizando descenso de gradiente estocástico (SGD) con una función de pérdida píxel a píxel, típicamente entropía cruzada o una variante ponderada. La función de energía se calcula mediante un soft-max píxel a píxel sobre el mapa de características final combinado con la pérdida de entropía cruzada. Los autores introdujeron un mapa de pesos que da mayor importancia a los píxeles cerca de los límites entre objetos en contacto, lo que ayuda a la red a aprender a separar instancias individuales en imágenes biomédicas. Este mapa de pesos se precalcula para cada imagen de entrenamiento y se utiliza para penalizar más fuertemente los errores en las regiones fronterizas.
La red utiliza aumento de datos para aumentar el tamaño efectivo del conjunto de entrenamiento, incluyendo deformaciones elásticas, rotaciones y variaciones de intensidad. Esto es particularmente importante en imágenes biomédicas donde los datos anotados son escasos. La implementación original utilizaba un momento alto (0.99) y una tasa de aprendizaje baja (0.00001) para asegurar una convergencia estable.
Aplicaciones en Imágenes Biomédicas
Hay muchas aplicaciones de U-Net en la segmentación de imágenes biomédicas, como la segmentación de diferentes órganos y sistemas corporales en tomografías computarizadas (TC) y resonancias magnéticas (RM). Casos específicos incluyen la segmentación de imágenes cerebrales (BRATS) y la segmentación de imágenes hepáticas (siliver07), así como la predicción de sitios de unión de proteínas. Las implementaciones de U-Net también han encontrado uso en las ciencias físicas, por ejemplo, en el análisis de micrografías de materiales.
Variaciones de U-Net también se han aplicado para la reconstrucción de imágenes médicas. Aquí hay algunas variantes y aplicaciones de U-Net a continuación:
- Regresión píxel a píxel usando U-Net y su aplicación en pansharpening.
- U-Net 3D: Aprendizaje de segmentación volumétrica densa a partir de anotación dispersa.
- TernausNet: U-Net con codificador VGG11 preentrenado en ImageNet para segmentación de imágenes.
- Traducción de imagen a imagen para estimar tinciones fluorescentes.
- En la predicción de sitios de unión de la estructura de proteínas.
Uso en Modelos de Difusión
La arquitectura U-Net también se ha empleado en modelos de difusión para la eliminación iterativa de ruido en imágenes. Esta tecnología subyace a muchos modelos modernos de generación de imágenes, como DALL-E, Midjourney y Stable Diffusion. En estos modelos, un U-Net se entrena para predecir el ruido añadido a una imagen en cada paso del proceso de difusión, permitiendo al modelo refinar gradualmente una entrada de ruido aleatorio en una imagen coherente. Las conexiones de salto en U-Net son particularmente beneficiosas aquí porque preservan detalles de alta frecuencia durante el proceso de eliminación de ruido.
Exploración para Modelos de Lenguaje
U-Net también se está explorando para modelos de lenguaje. La tokenización no es un paso separado, lo que permite al modelo entender más fácilmente la ortografía y vectorizar/tokenizar simultáneamente conceptos de nivel superior. Esta es un área emergente de investigación que aprovecha la capacidad de U-Net para capturar características multiescala, ofreciendo potencialmente una alternativa a la arquitectura Transformer (architecture) dominante en el procesamiento de lenguaje natural.
Historia y Legado
U-Net fue creado por Olaf Ronneberger, Philipp Fischer y Thomas Brox en 2015 y se informó en el artículo "U-Net: Convolutional Networks for Biomedical Image Segmentation". Es una mejora y desarrollo de FCN: Evan Shelhamer, Jonathan Long, Trevor Darrell (2014). "Fully convolutional networks for semantic segmentation". El artículo se ha convertido en uno de los más citados en el campo del aprendizaje profundo, y la arquitectura se ha adaptado para innumerables tareas más allá de las imágenes biomédicas.
El éxito de U-Net ha inspirado muchas variantes, incluyendo aquellas con diferentes codificadores (como VGG o ResNet), mecanismos de atención y procesamiento multiescala. Su influencia se extiende a otras áreas del aprendizaje automático, incluyendo IA generativa y visión por computadora. La arquitectura también es un punto de referencia común en desafíos de imágenes médicas y está ampliamente implementada en marcos como TensorFlow y PyTorch.
Implementaciones
Varias implementaciones de código abierto de U-Net están disponibles. Tensorflow Unet de J Akeret (2017) es una implementación popular. El código fuente original de U-Net está disponible en el grupo de Reconocimiento de Patrones y Procesamiento de Imágenes del Departamento de Ciencias de la Computación de la Universidad de Friburgo, Alemania. Estas implementaciones han facilitado la adopción generalizada de U-Net tanto en la investigación como en la industria.