UNet – Explicación paso a paso de la arquitectura

UNet [1] es una arquitectura de red neuronal profunda propuesta en el año 2015, destinada especialmente al ámbito de la imagen biomédica, ya que ha demostrado grandes resultados en este dominio.

A la hora de hablar de aprendizaje profundo (deep learning en inglés) y, más concretamente, de su aplicación en tareas de aprendizaje supervisado, es muy común hacer referencia a la resolución de problemas de clasificación.

Por ejemplo, una típica tarea sería la de clasificar una imagen en función de la raza de perro del ejemplar que aparezca en ella. En otras palabras, vamos a tener una imagen y una única etiqueta asociada a dicha imagen. En función del número de clases del problema, la salida del modelo podría no ser única.

Por ejemplo, si el problema planteado tuviera tres clases diferentes (Border Collie, Golden Retriever y Husky Siberiano), el modelo tendría tres salidas diferentes.

Ahora cambiemos completamente de tercio. Piensa en el dominio biomédico. Nosotros tenemos una imagen y queremos determinar si es patológica o no. Podemos resolver el mismo problema que antes. Tendríamos un modelo con una salida única que nos diga “patológica” o “no patológica”. Si tuviéramos distintos tipos de patologías, tendríamos tantas salidas como tipos de patologías más la salida correspondiente a la clase “no patológica”.

Aunque estas tareas también pueden ser interesantes bajo un punto de vista diagnóstico, resultaría muy útil poder tomar una imagen, procesarla y localizar de forma precisa las lesiones o, en general, las evidencias de la afectación de cierta patología en una imagen determinada.

Evidentemente, si la entrada al modelo de aprendizaje es una imagen de tamaño 512 x 512 y la salida es un vector de características con 3 clases (por tanto, de tamaño 1 x 3 ó 3 x 1, según si se quiere representar en columnas o en filas), dicha salida no va a tener ni mucho menos la suficiente resolución como para poder obtener un mapa donde localizar la extensión de una determinada patología.

Arquitecturas totalmente convolucionales

Lo ideal sería tener un modelo que fuera capaz de recibir la imagen de un tamaño determinado y devolver una imagen de ese mismo tamaño. Para ello, se han ideado diversas arquitecturas que se pueden encuadrar en el grupo de las redes totalmente convolucionales.

Si algo caracteriza a este tipo de redes es que soportan cualquier tipo de resolución a la entrada, ya que son capaces de adaptarse, por su composición interna, y ajustar el tamaño de la imagen de salida de la forma correspondiente.

Esto no quiere decir que las imágenes de entrada y las de salida tengan siempre la misma resolución. Esto depende, en gran medida, de cómo sea la resolución de la imagen de entrada y de cómo salgan los cálculos con los tamaños de kernel de las convoluciones.

Por lo general, es muy común que las arquitecturas totalmente convolucionales se ajusten bien a resoluciones de imagen de entrada que sean potencias de 2.

En el caso de tener una resolución arbitraria, sería necesario realizar un preprocesamiento de la imagen antes de que entre al modelo.

Por ejemplo, una manera sería cambiar la resolución de la imagen de entrada a un tamaño estándar. Otra forma podría ser añadiendo un padding a la imagen o realizando un recorte de la misma, hasta alcanzar las resoluciones deseadas y que se ajusten bien al problema.

Por eso, lo más recomendable siempre es trabajar con imágenes que tengan un ancho y alto iguales y con resoluciones que sean potencias de 2.

Características de una red totalmente convolucional

Las redes totalmente convolucionales se caracterizan por estar estructuradas en dos partes fundamentales. Concretamente, se suele decir que una red totalmente convolucional utiliza una arquitectura de encoder-decoder.

Esta arquitectura, como su propio nombre indica, está formada por un encoder y por un decoder.

El encoder es la primera parte de esta arquitectura y se encarga de reducir la dimensionalidad de la imagen de entrada. Básicamente, va reduciendo esa entrada hasta convertirla en un vector de características representativo de la misma. Esta reducción de la dimensionalidad se conoce como downsampling o submuestreo.

El decoder hace lo opuesto, aumentando la dimensionalidad de ese vector de características, aproximándose poco a poco a la resolución de imagen deseada a la salida. Este aumento de la dimensionalidad es lo que se conoce como upsampling o supermuestreo.

Una vez visto esto… pasemos a hablar de la UNet

Arquitectura de UNet (Ronneberger et al. [1])

Una vez comprendido cómo funciona una arquitectura totalmente convolucional, pasaremos a explicar paso a paso la UNet.

Antes de nada, es importante aclarar que se suele hablar de UNet como si fuese una arquitectura única que siempre se usa tal cual.

Lo cierto es que existen muchas variaciones, ya que existen trabajos más metodológicos que plantean nuevas arquitecturas de red basándose en la idea original de la UNet y, de igual modo, hay trabajos más aplicados que realizan pequeñas adaptaciones en la arquitectura para poder utilizarla en su caso particular.

Por tanto, nos vamos a basar en la implementación original para coger la idea más fundamental.

La arquitectura UNet original tiene 4 bloques de submuestreo (que, por tanto, se corresponden con el encoder), 4 bloques de supermuestreo (que, por tanto, se corresponden con el decoder) y un bloque intermedio al que se suele denominar como “bottleneck” o, en castellano, cuello de botella.

Adicionalmente, la arquitectura UNet incorpora el concepto de las skip connections, algo que en castellano se podría traducir como conexiones de salto. ¿Para qué sirve este concepto? Lo cierto es que el encoder va reduciendo poco a poco la dimensionalidad de la imagen de entrada, por lo que se pierde una gran cantidad de información espacial.

Para evitar esto, las skip connections llevan parte de esa información espacial del encoder directamente al decoder, lo cual puede ayudar al modelo en su proceso de aprendizaje.

Cabe destacar que, para esta arquitectura original de UNet, se presupone que la imagen de entrada está en escala de grises y que, por tanto, tiene un único canal de color.

Encoder de la UNet

Con respecto a las diferentes partes del encoder de la UNet, como bien se explica anteriormente, consta de 4 bloques de submuestreo. Cada uno de ellos tiene dos capas convolucionales con un kernel de 3 x 3.

Estas capas convolucionales, a su vez, utilizan una ReLU (rectificador lineal en castellano) como función de activación. Por último, cada bloque tendrá una capa de max pooling final de tamaño 2 x 2. Además, para el primer bloque de submuestreo, la primera capa convolucional el número de canales pasa de 1 a 64.

Para la segunda capa convolucional, el número de canales se mantendrá en 64. Para el resto de bloques de submuestreo, el número de canales se duplica de cada vez para la primera capa convolucional y se mantiene para la segunda (alcanzando, de este modo, 512 canales para el último bloque del encoder).

Del mismo modo, al final del encoder, el tamaño de la imagen se habrá reducido hasta una resolución de 32 x 32.

Conceptualmente, es importante entender que las capas de max pooling son las que llevan a cabo esa reducción de la dimensionalidad. Por su parte, las capas convolucionales son operaciones que permiten filtrar las imágenes y extraer características relevantes de las mismas.

Es por ello que las capas convolucionales, a menudo, aumentan el número de canales. Cada uno de esos canales será, por tanto, un filtro distinto.

Una vez en el bottleneck, se aplican dos pasos convolucionales más (con las mismas características, es decir, con un kernel de 3 x 3 y una función de rectificación lineal) hasta llegar a una resolución de 28 x 28.

Decoder

Una vez pasado el encoder y el bottleneck, comienzan las etapas de supermuestreo del decoder.

En estas fases, lo que conseguimos es aumentar la resolución de forma sucesiva hasta alcanzar el tamaño deseado. Esto se logra esencialmente gracias a las convoluciones transpuestas, un método de supermuestreo alternativo a las clásicas interpolaciones, que presentan la desventaja de pixelar y emborronar las imágenes cuando el reescalado es demasiado cuantioso.

De este modo, por tanto, el modelo es capaz de aprender cómo debe supermuestrear la imagen para tener la resolución adecuada y que la salida de la red tenga una calidad satisfactoria. Del mismo modo, garantiza que el modelo persiga el objetivo de representar lo que se desea representar (en este caso, la segmentación de una imagen).

Dentro del decoder, cada bloque de supermuestreo tiene dos capas convolucionales de tamaño 3 x 3, cada una de ellas con su correspondiente rectificador lineal. Estas capas convolucionales permiten ir reduciendo poco a poco el número de canales de las imágenes, hasta alcanzar un número de canales equivalente al número de clases del problema.

El segundo detalle importante es que, después de encadenar esas dos capas de convoluciones, los bloques de supermuestreo incluyen al final una capa de convolución transpuesta con un kernel de tamaño 2 x 2. Esta capa final de cada bloque es la que se encarga de llevar a cabo el aumento de la dimensionalidad siendo, dicho a grandes rasgos, la operación opuesta a las capas de max pooling del encoder.

Además, se implementan skip connections que unen cada bloque de submuestreo con su correspondiente bloque de supermuestreo.

Tomando el primer bloque de submuestreo como referencia, tras los dos pasos convolucionales, la imagen de salida tendrá un número de canales de 64.

Para concatenar parte de la información espacial que se pierde en el submuestreo con la imagen que se va agrandando en el supermuestreo, se toman los 64 canales de ese bloque de submuestreo y se concatenan con los 64 que se obtienen de modo “natural” en esa fase del supermuestreo.

Finalmente, y dado que la salida de la red es una imagen con 64 canales, es necesario reducir esto al número de clases del problema. Esto se consigue por medio de una capa convolucional final de 1 x 1, utilizando la función de activación deseada, que generalmente suele ser una sigmoidal, una tangente hiperbólica o una función lineal o cualquier otra función que se considere adecuada para resolver un problema de clasificación.

En la arquitectura original, por ejemplo, la imagen de salida tiene 2 canales. De esta forma, entrenando un modelo como este, tendremos un sistema capaz de recibir una imagen a la entrada y dar su correspondiente segmentación a la salida, proporcionando un valor de clasificación por cada clase en cada uno de los píxeles de la imagen.

Cabe destacar que, para esta arquitectura concreta de UNet y para Para evitar esto, como hemos indicado en otras partes de este post, sería necesario realizar algún tipo de preprocesamiento para ajustar el tamaño de la imagen de salida.

Referencias

[1] Ronneberger, O., Fischer, P., & Brox, T. (2015, October). U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention (pp. 234-241). Springer, Cham.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

Create your website with WordPress.com
Get started
%d bloggers like this: