TLDR: Will a u-net pretrained as an autoencoder be able to learn a latent representation of the data if the encoder weights are frozen (can't game the system and pass forward the unmodified image)? Literature on this technique?

I wrote a u-net (Notebook: https://www.kaggle.com/code/gordonbchen/gi-polyp-segmentation) that uses MobileNetV2 as its encoder backbone to segment images of colon polyps from the KVASIR dataset (https://paperswithcode.com/dataset/kvasir).

I froze the weights for the MobileNetV2 backbone and built a u-net autoencoder that takes in images of the colon polyps without segmentation masks. I wanted to pretrain the u-net because images with high-quality segmentation masks are hard to come by (1000 in the KVASIR seg dataset), while images of the same thing without the masks are more accessible (4000 in the KVASIR dataset).

I trained the autoencoder and got a pretty exact mapping from the input to output image, which is really suspicious, and makes me think that the u-net isn't learning a latent representation of colon polyp images, but has just figured out a way to pass the unmodified image through the skip (concat) connections.

See auto-encoder output below (top row = input, middle row = target output (same as input b/c autoencoder), bottom row = model prediction). U-Net autoencoder predictions

I then trained the segmentation u-net, which just has a different output layer than the autoencoder, and ended up getting a better result (reaches a lower validation loss, less epochs needed training on segmentation dataset, and much more stable training), than when I trained the u-net without autoencoder pretraining.

U-net autoencoder training loss. enter image description here

Segmentation u-net training loss after autoencoder pretraining. Segmentation u-net training loss after autoencoder pretraining

Segmentation u-net training loss (no autoencoder pretraining). Segmentation u-net training loss (no autoencoder pretraining)

Essentially I'm wondering if there is any literature on pretraining u-nets as autoencoders or if this is a new approach, whether or not this approach is viable (if pretraining with an autoencoder will actually teach the model something, and not just teach it to propagate the unchanged image forward), and if the benefits I saw from doing so are explained and expected. Since the encoder (MobileNetV2) is frozen, I don't think the network can "game the system" as much as if I trained a u-net autoencoder from scratch.

Final outputs (with autoencoder pretraining) Final outputs (with autoencoder pretraining)

Edit 1/16/2024: also posted question on Kaggle: https://www.kaggle.com/discussions/questions-and-answers/468463

  • $\begingroup$ If you want to do unsupervised/self-supervised pretraining of a unet model, the typical setup would be to create a paired dataset where the inputs are a degraded version of the image, and the outputs are the clean image. The self-supervised training task is to undo the image degradation. $\endgroup$
    – Karl
    Commented Jan 17 at 1:42


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.