Mind2Mind : transfer learning for GANs
Jean-Baptiste Gouray • Yaël Frégier
Official repository of the paper
This repository contains a Mind2Mind transfer module. We have added it to a fork of the ALAE repository. We have kept from this fork only the modules essential for running Mind2Mind. If you need the full capacities of ALAE, add our module to the original ALAE repository.
Google Drive folder with models and qualitative results
Mind2Mind
Transfer Learning for GANs
Abstract: Training generative adversarial networks (GANs) on high quality (HQ) images involves important computing resources. This requirement represents a bottleneck for the development of applications of GANs. We propose a transfer learning technique for GANs that significantly reduces training time. Our approach consists of freezing the low-level layers of both the critic and generator of the original GAN. We assume an auto-encoder constraint in order to ensure the compatibility of the internal representations of the critic and the generator. This assumption explains the gain in training time as it enables us to bypass the low-level layers during the forward and backward passes. We compare our method to baselines and observe a significant acceleration of the training. It can reach two orders of magnitude on HQ datasets when compared with StyleGAN. We prove rigorously, within the framework of optimal transport, a theorem ensuring the convergence of the learning of the transferred GAN. We moreover provide a precise bound for the convergence of the training in terms of the distance between the source and target dataset.
Repository organization
To run the scripts, you will need to have a CUDA capable GPU, PyTorch >= v1.3.1 and cuda/cuDNN drivers installed. Install the required packages:
pip install -r requirements.txt
Running scripts
The code in the repository is organized in such a way that all scripts must be run from the root of the repository. If you use an IDE (e.g. PyCharm or Visual Studio Code), just set Working Directory to point to the root of the repository.
If you want to run from the command line, then you also need to set PYTHONPATH variable to point to the root of the repository.
For example, let's say we've cloned repository to ~/ALAE directory, then do:
$ cd ~/ALAE
$ export PYTHONPATH=$PYTHONPATH:$(pwd)
Now you can run scripts as follows:
$ python module_mind/generate_images.py
Repository structure
Path | Description |
---|---|
ALAE | Repository root folder |
├ configs | Folder with yaml config files. |
│ └ ffhq.yaml | Config file for FFHQ dataset at 1024x1024 resolution. |
├ module_mind | Folder with Mind2Mind module. |
│ ├ data_loader.py | Class to define loaders for encoded data. |
│ ├ download_mindGAN.py | Script to download a pre-trained MindGan. |
│ ├ generate_images.py | Script to generate images from the MindGAN. |
│ ├ model.py | MindGAN model. |
│ ├ prepare_data.py | Script to download celebaHQ and encode the data. |
│ ├ train.py | Script to train the MindGAN on CelebaHQ from the ALAE autoencoder trained FFHQ. |
│ └ trainer.py | Class for handling training loops. |
├ checkpointer.py | Module for saving/restoring model weights, optimizer state and loss history. |
├ defaults.py | Definition for config variables with default values. |
├ losses.py | Defintions of the loss fonctions. |
├ lreq.py | Custom Linear , Conv2d and ConvTranspose2d modules for learning rate equalization. |
├ model.py | Module with high-level model definition. |
├ net.py | Definition of all network blocks for multiple architectures. |
├ registry.py | Registry of network blocks for selecting from config file. |
├ random_choice.png | Sample of images (for this readme). |
├ requirements.txt | List of python modules needed. |
└ utils.py | Decorator for async call, decorator for caching, registry for network blocks. |
Configs
In ALAE, you can specify which yaml config yacs will use. However, our Mind2Mind module only accepts for the moment the ffhq
config. Since it is the default config for ALAE you do not have anything to do. However, if you know the use of the -c
parameter from ALAE, do not try to use it here to choose another config.
Datasets
You must prepare the data with the command:
$ python module_mind/prepare_data.py
Pre-trained models
To download pre-trained models run:
python module_mind/download_mindGAN.py
Generating figures
To make generation figure run:
python module_mind/generate_images.py
By default, it will generate one batch of 4 images. If you want to modify the numbers of batches and images, you have to modify the lines 19-20 in generate_images
. In particular if your system runs out of memory, you will need to lower the number of images par batch and restart the kernel. The generated samples can be found in the folder module_mind/images_generated/mind2mid
.
Training
To run training:
python module_mind/train.py
We have only tested our Mind2Mind module on a single GPU.
You might need to adjust the batch size in the config file depending on the memory size of the GPU.
Computation of FID
To compute the fid score, you need to download the module pytorch-fid and run the command :
python fid_score $ALAE_PATH/module_mind/Dataset/Celeba-HQ/data1024x1024 $ALAE_PATH/module_mind/images_generated/mind2mind/
where $ALAE_PATH
is the directory in which ALAE is located.
Our settings description
{\bf Datasets}. We have tested our algorithm at resolution
Description of hyperparameters
{\bf At resolution 28$\times $28}. The encoder
{\bf At resolution 1024$\times