Image Colorization: U-Net Noob Solution

Image Colorization implementation in PyTorch using U-Net: from Data Preprocessing to Model Training

Author's Note!

The model presented here was my final project for my ML training course. The whole code was written by me except adopted implementation of moving average from ZELDA framework by Michael Konstantinov and model summary for PyTorch implemented by Jong Wook Kim. The sources I used for creating the project you can find in the model notebook or right here after the main text body. I faced a lot of tricky issues implementing the project (they are discussed below too), so I hope this code will help someone. Please, feel free to use the code for your education. Also I'm open to comments and suggestions. If this repo helps you, please, let me know ;)

Brief overview

Notebook Google Colab Pro

Datasets CIFAR-100, Dogs-vs-Cats

Architecture U-Net

Framework PyTorch

Optimizer Adam

Loss MSE

Notebook

I used Google Colaboratory as that was easy and comfortable way for tutor and students to share projects between each other. Google Colab gives you 12,5 GB of RAM for free and it's enough if you work with datasets like CIFAR-10 or CIFAR-100. If you want to train your model with a heavier dataset (e.g., Dogs-vs-Cats, CelebA), you need at least about 20 GB of RAM to feel yourself free. Colab Pro provides 25 GB.

Architecture - U-Net

U-Net is a fully convolutional network. It has a u-shape architecture that actually is an encoder-decoder type. U-Net has 31,042,434 parameters, so the training holds from few hours to half of the day depending on the dataset you use. So be patient.

I trained 3 models

Framework - PyTorch

I lust like it. The end.

Optimizer - Adam

The only true option. SGD and RMSProp slow down training a lot. All the best Image Colorization models use Adam.

Loss - MSE

"Good default choice" for Image Colorization, but actually it's the worst option you can ever imagine. MSE is too straightforward for such comlicated regression task as Image Colorization. It makes model to learn the most popular colours, so the colorized images are usually brownish and greenish. The best choice is custom Loss (weighted Loss) that gives an andvantage to rare colours. Unfourtunately, I didn't have time to try custom Loss for this model. This project is a good exmaple why you should NOT use MSE for Image Colorization.

Datasets

Dataset also has impact on sepia effect of colorized images. Colorful datasets are the best choices as you model can learn more colours.