Segmentación de melanoma con U-Net implementadas desde cero en PyTorch

Problema

Delimitar con precisión el borde de una lesión de piel es el primer paso para analizarla: la asimetría y la irregularidad del borde son señales clásicas de melanoma. El objetivo es un modelo que, dada una imagen dermatoscópica, prediga la máscara de la lesión píxel a píxel (lesión o fondo).

El conjunto es ISIC Melanoma Segmentation: 2.750 imágenes con su máscara anotada. La pregunta del proyecto es si compensa partir de un codificador preentrenado en ImageNet, un dominio muy distinto, frente a entrenar toda la red desde cero.

Mi contribución

Implementé desde cero en PyTorch una U-Net clásica: cinco bloques convolucionales con filtros de 64 a 1024, decodificador simétrico con convoluciones transpuestas y conexiones de salto en cada nivel (31 M de parámetros, todos entrenables).
Implementé una U-Net con codificador ResNet-50 preentrenado y congelado, con conexiones de salto desde conv1, layer1, layer2 y layer3, y un decodificador propio de cinco bloques. Entrena 26,5 M de sus 50 M de parámetros.
Escribí el pipeline completo: carga y preprocesado, entrenamiento con early stopping, selección del umbral y evaluación.

Enfoque técnico

Preprocesado512 × 512 px, escala de grises, normalización y máscara binarizada
EntrenamientoAdam (1e-4) y BCE con logits, hasta 50 épocasearly stopping: AUC val.
Umbral óptimoCurva precisión-recall en validación; se elige el umbral de máximo F1
Evaluación en testAUC-ROC por píxel e IoU (índice de Jaccard)

Pipeline común a los dos modelos. Partición aleatoria 70/10/20: 1.924 imágenes de entrenamiento, 275 de validación y 551 de test.

Las imágenes pasan a escala de grises, así que la primera capa de ResNet-50, pensada para RGB, no encaja. En lugar de inicializarla al azar, promedié sus pesos de los tres canales: conserva lo aprendido sobre bordes y texturas y acepta un solo canal.

# scripts/models/UNet_ResNet50_Encoder.py (fragmento)
resnet50 = models.resnet50(pretrained=True)
self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
# De RGB a escala de grises: media de los pesos de los tres canales
self.conv1.weight.data = resnet50.conv1.weight.data.mean(dim=1, keepdim=True)

# Codificador congelado: solo se entrena el decodificador
encoder = [self.conv1, self.bn1,
           self.layer1, self.layer2, self.layer3, self.layer4]
for layer in encoder:
    for param in layer.parameters():
        param.requires_grad = False

Dos decisiones del protocolo de evaluación:

AUC por muestreo. Calcular la AUC con todos los píxeles de validación no cabe en memoria; se calcula con uno de cada diez, que basta para guiar el early stopping.
Umbral sin fugas. El umbral que convierte probabilidades en máscara se elige en validación (máximo F1 en la curva precisión-recall) y después se aplica tal cual en test. Elegirlo en test inflaría el resultado.

Resultados

Modelo	Parámetros entrenables	AUC-ROC (test)	IoU (test)	Umbral
U-Net desde cero	31,0 M	0,9616	0,7236	0,42
U-Net con ResNet-50 congelado	26,5 M	0,9769	0,7524	0,34

El codificador preentrenado gana en las dos métricas (+1,5 puntos de AUC y +2,9 de IoU) aunque entrena menos parámetros. Las características de ImageNet, aun siendo de otro dominio, capturan bordes, texturas y formas que aceleran la convergencia, y congelarlas actúa como regularización con un conjunto de entrenamiento de menos de 2.000 imágenes.

Dos casos de test en tres columnas: imagen original, máscara real y máscara predicha por la U-Net con ResNet-50. En el segundo caso la predicción sigue el borde irregular de la lesión con más detalle que la máscara real, que es un polígono. — U-Net con ResNet-50 en test: imagen, máscara real y predicción. La máscara real es un polígono aproximado; la predicción sigue el borde irregular de la lesión.

Dos casos de test de la U-Net entrenada desde cero: imagen original, máscara real y máscara predicha. Las predicciones cubren bien las lesiones grandes, con algunos huecos y ruido en los bordes. — U-Net desde cero en test. Acierta la forma general de lesiones grandes, con más huecos y ruido en los bordes.

Curvas de entrenamiento de la U-Net con ResNet-50. Izquierda: pérdida de entrenamiento y validación, que bajan de 0,33 a 0,14. Derecha: AUC de validación, que sube de 0,955 a un máximo de 0,973 en la época 26. — Entrenamiento de la U-Net con ResNet-50. La AUC de validación alcanza su máximo (0,973) en la época 26 y el early stopping detiene el entrenamiento cinco épocas después.

Evidencias

Repositorio con las dos arquitecturas (scripts/models/) y el notebook de entrenamiento con todas las salidas.

Aplicación profesional

Segmentación en imagen médica (lesiones, órganos, tumores) y en otros dominios: defectos en inspección industrial, parcelas en imagen satelital.
Transfer learning cuando hay pocos datos etiquetados: reutilizar un codificador preentrenado y entrenar solo lo específico del problema.
Protocolos de evaluación sin fugas, con los umbrales y decisiones de modelo tomados siempre fuera del conjunto de test.