In the first step, the generator creates a new image. Then the discriminator verifies whether the generated photo is real or not. At this point, GAN calculates two losses, for generator and discriminator. The generator learns how good was his image and the discriminator learns how good was his verification.
ESRGAN is the pre-trained model with vgg19 weights. Do Not Reinvent The Wheel. (Artist DALL-E)
How to prepare ESRGAN for your purpose
Firstly you have to upload your dataset to Google Drive. In this tutorial, I will use the dataset from Kaggle called CalebA. This dataset contains over 200k images of celebrities' faces. The resolution is 218x178 with 3 channels.
If you want to use the same dataset, I recommend uploading only 10 000 images.
Nevertheless, you can use whatever you want!
ESRGAN needs GPU with a large memory capacity, so in the tutorial, I use Google Colab. In which you can change Runtime type, click Runtime and choose GPU in Hardware accelerator.
Now we need to clone the repository with implemented ESRGAN
and install requirements.
Then, we can connect Google Drive with Google Colab by command:
I use the patool lib to extract files from my rar to Google Colab.
Important outdir must be "/content/PyTorch-GAN/data"
A good practice is to have a testing dataset. The data was not used to train a model. To create testing set we can transfer some data from data/name_of_your_folder to data/test folder.
Firstly we have to create a new folder
After that, we move a few images to test folder.
If the uploaded count of images is too big. You can use:
Let's train the neural network!
To train the ESRGAN model we have to write a command with arguments:
Available arguments:
--dataset_name name of your folder in /content/PyTorch-GAN/data
--n_epochs number of epochs (default 200)
--hr_heightheight of output (default 256)
--hr_width width of output (default 256)
--channeslchannels of input (default 3)
--checkpoint_interval I recommend setting it as 250 if you have a small dataset. (default 5000)
The rest arguments you can find HERE.
Inside the folder '/content/PyTorch-GAN/implementations/esrgan/images/training you can find the saved images from the training process.
To test our model we need an image. We use the images which we have transferred before to test folder.
We can run the model with the command:
Available arguments:
--image_path name of your image in /content/PyTorch-GAN/data/test/0.jpg
--checkpoint_model name of your trained generator /content/PyTorch-GAN/implementations/esrgan/saved_models/generator_X.pth
X replace with a number of the last trained epoch.
The generated image you can find in /content/PyTorch-GAN/implementations/esrgan/images/outputs/
If you want to copy the image to your google cloud write
The idea of GAN models is amazing, two neural networks play together to be better in every epoch -Use a neural network to teach a neural network. The GAN models have more and more types, one of them is used in the tutorial ESRGAN. These neural networks are prepared to work with images and make them better. This is a powerful model, but it also needs a lot of computing power. Below you can see the results after 5 epochs and 10 000 training data. They are not ideal, but if you run your training process for much more epochs the results will be awesome! Just try it!