train.py is the generic training script . It serves many models (optional models like: pix2pix, cyclegan, colorization) and different datasets (optional dataset modes include: aligned, unaligned, single, colorization).
test.py is a generic test script . If you have trained your own model with train.py, you can use the test script to test the model. The script will load the saved model from checkpoints_dir and save the result in results_dir.
The data directory includes all modules related to data loading and preprocessing . If you want to add a custom dataset class named dummy, you need to create a dummy_dataset.py file and define a subclass (subclass) Dummy Dataset that inherits from Base Dataset. Four functions need to be implemented: 1) __init__ (initialization class, first call BaseDataset.__init__(self, opt)), 2) __len__ (return data set size), 3) __getitem__ (get data points), 4) modify_commandline_options (add Dataset-specific options, set default options). A dataset class can be used by specifying --dataset_mode. Each file is explained in detail below.
__init__.py implements the interface of the data package and training and testing script files, train.py and test.py call data from data import create_dataset and dataset = create_dataset(opt), and create a dataset with the given option opt.
base_dataset.py implements the theoretical base class ( ABC ) for datasets. It also includes common transformation functions (such as get_transform, __scale_width), which can be used in later subclasses.
image_folder.py implements an image folder class. We modified the official PyTorch image folder code so that this class can load images from the current directory and its subdirectories.
template_dataset.py provides a dataset template with detailed documentation. Check out this file if you plan to implement your own dataset.
aligned_dataset.py contains a dataset class that can load image pairs. It assumes a single image directory /path/to/data/train that contains image pairs of the form {A,B}. When testing, you need to prepare the directory /path/to/data/test as test data.
unaligned_dataset.py contains a dataset class that can load unaligned/unpaired datasets. It assumes that two directories host training images from domain A /path/to/data/trainA and domain B /path/to/data/trainB respectively. Models can be trained with the dataset flag --dataroot /path/to/data . Likewise, two directories /path/to/data/testA and /path/to/data/testB need to be prepared during testing.
single_dataset.py contains a dataset class that can load a set of single images by specifying the path --dataroot /path/to/data. It can generate unilateral CycleGAN results with an optional test model.
colorization_dataset.py implements a dataset class that can load a set of RGB natural images, and convert the RGB format into (L, ab) pairs in the Lab color space. It is required by the Pix2Pix based shading model (--model colorization).
The models directory contains modules related to objective functions, optimization, and network architectures . To add a custom dataset class named dummy, you need to create a dummy_model.py file and define a subclass (subclass) Dummy Model that inherits Base Model. The following functions need to be implemented: 1) __init__ (initialize the class, first call BaseModel.__init__(self, opt)), 2) set__input (unpack the data from the dataset and perform preprocessing), 3) forward (generate intermediate results), 4) optimize_parameters (calculate loss, gradient, update network weights). 5) Modify_commandline_options can also be added (set default options by adding model-specific options). Model classes can be used by specifying the flag --model dummy. Each file is explained in detail below.
__init__.py implements the interface between the model package and the training and testing script files, train.py and test.py use from models import create_model and model = create_model(opt) to create a model with the given option opt. Initialize the model with model.setup(opt).
base_model.py implements the theoretical base class ( ABC ) for the model. It also includes common helper functions (such as: setup, test, update_learning_rate, save_networks, load_networks), which will be used in subclasses later.
template_model.py provides model templates with detailed documentation. If you want to implement your own model, you can check this file.
pix2pix_model.py implements the Pix2Pix model for learning an input-to-output mapping of pairs of images. Model training requires --dataset_mode aligned dataset. By default it uses --netG unet256 network, --netD basic discriminator (PatchGAN) and --gan_model vanillagan loss (standard cross-entropy objective).
colorization_model.py implements a subclass of Pix2PixModel for image colorization (black and white image to color image). Model training requires -dataset_model colorization dataset. It trains a Pix2Pix model that maps from the L channel to the ab channel in the Lab color space. By default, the colorization dataset will automatically set --input_nc 1 and --output_nc 2.
cycle_gan_model.py implements the cycle gan model for learning unpaired image-to-image translation. Model training requires --dataset_mode unaligned dataset. By default it uses --netG resnet_9blocks generator, --netD basic discriminator and least squares Gans objective ( --gan_model lsgan ).
The networks.py module implements network architectures (generators and discriminators), as well as normalization layers, initialization methods, optimization schedulers (i.e. learning rate policies), and GAN objective functions (vanilla, lsgan, wgangp).
test_model.py implements the model used to generate one-sided Cyclegan results. This model will automatically set --dataset_mode single which only loads images from one collection. (--model cycle_gan can be tested in both directions)
The options directory includes options modules : training options, testing options, and basic options (for training and testing). TrainOptions and TestOptions are two subclasses of BaseOptions. They will recall the options defined in BaseOptions.
__init__.py causes Python to treat the options directory as a containing package.
base_options.py includes options used in training and testing. It also implements some helper functions such as parsing, printing and saving options. It also collects other options defined in the modify_commandline_options function in dataset classes and model classes.
train_options.py only includes options used during training.
test_options.py only includes options used during testing.
The util directory contains a collection of various useful helper functions .
__init__.py causes Python to treat the util directory as an included package.
get_data.py provides a Python script for downloading the CycleGAN and pix2pix datasets. For example: download_pix2pix_model.sh or download_cyclegan_model.sh .
html.py is used to save images into a separate HTML file. It includes functions such as: add_header (add text headers to HTML files), add_images (add a line of images to HTML files), save (save HTML to disk). It is based on the Python library dominate, a python library for creating and manipulating HTML documents with the DOM API.
image_pool.py is an image buffer that stores previously generated images. This buffer allows us to update the discriminator with previously generated images instead of the most recent generator. The paper discusses this initial idea. The size of the buffer is controlled by --pool_size.
visualizer.py can display/save images and print/save log information. It uses the Python library visdom for display and the Python library dominate (wrapped in HTML) to create HTML files with images.
util.py contains helper functions such as tensor2im (converts tensor arrays to numpy image arrays), diagnose_network (computes and prints the mean of the mean absolute value of gradients), mkdirs (creates multiple directories).
## Overview of Code Structure
To help users better understand and use our codebase, we briefly overview the functionality and implementation of each package and each module. Please see the documentation in each file for more details. If you have questions, you may find useful information in [training/test tips](tips.md) and [frequently asked questions](qa.md).
[train.py](../train.py) is a general-purpose training script. It works for various models (with option `--model`: e.g., `pix2pix`, `cyclegan`, `colorization`) and different datasets (with option `--dataset_mode`: e.g., `aligned`, `unaligned`, `single`, `colorization`). See the main [README](.../README.md) and [training/test tips](tips.md) for more details.
[test.py](../test.py) is a general-purpose test script. Once you have trained your model with `train.py`, you can use this script to test the model. It will load a saved model from `--checkpoints_dir` and save the results to `--results_dir`. See the main [README](.../README.md) and [training/test tips](tips.md) for more details.
[data](../data) directory contains all the modules related to data loading and preprocessing. To add a custom dataset class called `dummy`, you need to add a file called `dummy_dataset.py` and define a subclass `DummyDataset` inherited from `BaseDataset`. You need to implement four functions: `__init__` (initialize the class, you need to first call `BaseDataset.__init__(self, opt)`), `__len__` (return the size of dataset), `__getitem__` (get a data point), and optionally `modify_commandline_options` (add dataset-specific options and set default options). Now you can use the dataset class by specifying flag `--dataset_mode dummy`. See our template dataset [class](../data/template_dataset.py) for an example. Below we explain each file in details.
* [\_\_init\_\_.py](../data/__init__.py) implements the interface between this package and training and test scripts. `train.py` and `test.py` call `from data import create_dataset` and `dataset = create_dataset(opt)` to create a dataset given the option `opt`.
* [base_dataset.py](../data/base_dataset.py) implements an abstract base class ([ABC](https://docs.python.org/3/library/abc.html)) for datasets. It also includes common transformation functions (e.g., `get_transform`, `__scale_width`), which can be later used in subclasses.
* [image_folder.py](../data/image_folder.py) implements an image folder class. We modify the official PyTorch image folder [code](https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py) so that this class can load images from both the current directory and its subdirectories.
* [template_dataset.py](../data/template_dataset.py) provides a dataset template with detailed documentation. Check out this file if you plan to implement your own dataset.
* [aligned_dataset.py](../data/aligned_dataset.py) includes a dataset class that can load image pairs. It assumes a single image directory `/path/to/data/train`, which contains image pairs in the form of {A,B}. See [here](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md#prepare-your-own-datasets-for-pix2pix) on how to prepare aligned datasets. During test time, you need to prepare a directory `/path/to/data/test` as test data.
* [unaligned_dataset.py](../data/unaligned_dataset.py) includes a dataset class that can load unaligned/unpaired datasets. It assumes that two directories to host training images from domain A `/path/to/data/trainA` and from domain B `/path/to/data/trainB` respectively. Then you can train the model with the dataset flag `--dataroot /path/to/data`. Similarly, you need to prepare two directories `/path/to/data/testA` and `/path/to/data/testB` during test time.
* [single_dataset.py](../data/single_dataset.py) includes a dataset class that can load a set of single images specified by the path `--dataroot /path/to/data`. It can be used for generating CycleGAN results only for one side with the model option `-model test`.
* [colorization_dataset.py](../data/colorization_dataset.py) implements a dataset class that can load a set of nature images in RGB, and convert RGB format into (L, ab) pairs in [Lab](https://en.wikipedia.org/wiki/CIELAB_color_space) color space. It is required by pix2pix-based colorization model (`--model colorization`).
[models](../models) directory contains modules related to objective functions, optimizations, and network architectures. To add a custom model class called `dummy`, you need to add a file called `dummy_model.py` and define a subclass `DummyModel` inherited from `BaseModel`. You need to implement four functions: `__init__` (initialize the class; you need to first call `BaseModel.__init__(self, opt)`), `set_input` (unpack data from dataset and apply preprocessing), `forward` (generate intermediate results), `optimize_parameters` (calculate loss, gradients, and update network weights), and optionally `modify_commandline_options` (add model-specific options and set default options). Now you can use the model class by specifying flag `--model dummy`. See our template model [class](../models/template_model.py) for an example. Below we explain each file in details.
* [\_\_init\_\_.py](../models/__init__.py) implements the interface between this package and training and test scripts. `train.py` and `test.py` call `from models import create_model` and `model = create_model(opt)` to create a model given the option `opt`. You also need to call `model.setup(opt)` to properly initialize the model.
* [base_model.py](../models/base_model.py) implements an abstract base class ([ABC](https://docs.python.org/3/library/abc.html)) for models. It also includes commonly used helper functions (e.g., `setup`, `test`, `update_learning_rate`, `save_networks`, `load_networks`), which can be later used in subclasses.
* [template_model.py](../models/template_model.py) provides a model template with detailed documentation. Check out this file if you plan to implement your own model.
* [pix2pix_model.py](../models/pix2pix_model.py) implements the pix2pix [model](https://phillipi.github.io/pix2pix/), for learning a mapping from input images to output images given paired data. The model training requires `--dataset_mode aligned` dataset. By default, it uses a `--netG unet256` [U-Net](https://arxiv.org/pdf/1505.04597.pdf) generator, a `--netD basic` discriminator (PatchGAN), and a `--gan_mode vanilla` GAN loss (standard cross-entropy objective).
* [colorization_model.py](../models/colorization_model.py) implements a subclass of `Pix2PixModel` for image colorization (black & white image to colorful image). The model training requires `-dataset_model colorization` dataset. It trains a pix2pix model, mapping from L channel to ab channels in [Lab](https://en.wikipedia.org/wiki/CIELAB_color_space) color space. By default, the `colorization` dataset will automatically set `--input_nc 1` and `--output_nc 2`.
* [cycle_gan_model.py](../models/cycle_gan_model.py) implements the CycleGAN [model](https://junyanz.github.io/CycleGAN/), for learning image-to-image translation without paired data. The model training requires `--dataset_mode unaligned` dataset. By default, it uses a `--netG resnet_9blocks` ResNet generator, a `--netD basic` discriminator (PatchGAN introduced by pix2pix), and a least-square GANs [objective](https://arxiv.org/abs/1611.04076) (`--gan_mode lsgan`).
* [networks.py](../models/networks.py) module implements network architectures (both generators and discriminators), as well as normalization layers, initialization methods, optimization scheduler (i.e., learning rate policy), and GAN objective function (`vanilla`, `lsgan`, `wgangp`).
* [test_model.py](../models/test_model.py) implements a model that can be used to generate CycleGAN results for only one direction. This model will automatically set `--dataset_mode single`, which only loads the images from one set. See the test [instruction](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix#apply-a-pre-trained-model-cyclegan) for more details.
[options](../options) directory includes our option modules: training options, test options, and basic options (used in both training and test). `TrainOptions` and `TestOptions` are both subclasses of `BaseOptions`. They will reuse the options defined in `BaseOptions`.
* [\_\_init\_\_.py](../options/__init__.py) is required to make Python treat the directory `options` as containing packages,
* [base_options.py](../options/base_options.py) includes options that are used in both training and test. It also implements a few helper functions such as parsing, printing, and saving the options. It also gathers additional options defined in `modify_commandline_options` functions in both dataset class and model class.
* [train_options.py](../options/train_options.py) includes options that are only used during training time.
* [test_options.py](../options/test_options.py) includes options that are only used during test time.
[util](../util) directory includes a miscellaneous collection of useful helper functions.
* [\_\_init\_\_.py](../util/__init__.py) is required to make Python treat the directory `util` as containing packages,
* [get_data.py](../util/get_data.py) provides a Python script for downloading CycleGAN and pix2pix datasets. Alternatively, You can also use bash scripts such as [download_pix2pix_model.sh](../scripts/download_pix2pix_model.sh) and [download_cyclegan_model.sh](../scripts/download_cyclegan_model.sh).
* [html.py](../util/html.py) implements a module that saves images into a single HTML file. It consists of functions such as `add_header` (add a text header to the HTML file), `add_images` (add a row of images to the HTML file), `save` (save the HTML to the disk). It is based on Python library `dominate`, a Python library for creating and manipulating HTML documents using a DOM API.
* [image_pool.py](../util/image_pool.py) implements an image buffer that stores previously generated images. This buffer enables us to update discriminators using a history of generated images rather than the ones produced by the latest generators. The original idea was discussed in this [paper](http://openaccess.thecvf.com/content_cvpr_2017/papers/Shrivastava_Learning_From_Simulated_CVPR_2017_paper.pdf). The size of the buffer is controlled by the flag `--pool_size`.
* [visualizer.py](../util/visualizer.py) includes several functions that can display/save images and print/save logging information. It uses a Python library `visdom` for display and a Python library `dominate` (wrapped in `HTML`) for creating HTML files with images.
* [util.py](../util/util.py) consists of simple helper functions such as `tensor2im` (convert a tensor array to a numpy image array), `diagnose_network` (calculate and print the mean of average absolute value of gradients), and `mkdirs` (create multiple directories).
The code composition of CycleGAN-Jianshu (jianshu.com) 
https://www.jianshu.com/p/c77d69f6aa18 Relevant explanations: