add comments: models

Author: junyanz

data/__init__.py

@@ -1,11 +1,12 @@
 """This package includes 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.
+ 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, first call BaseDataset.__init__(self, opt))
-    -- <__len__> (return the size of dataset)
-    -- <__getitem__>　(get a data point)
-    -- (optionally) <modify_commandline_options> (add dataset-specific options and set default options).
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point from data loader.
+    -- <modify_commandline_options>:    (optionally) 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 'template_dataset.py' for more details.
 """
@@ -15,7 +16,7 @@
 
 
 def find_dataset_using_name(dataset_name):
-    """Import the module "data/[datasetname]_dataset.py" given the option '--dataset_mode [datasetname].
+    """Import the module "data/[dataset_name]_dataset.py".
 
     In the file, the class called DatasetNameDataset() will
     be instantiated. It has to be a subclass of BaseDataset,
@@ -44,14 +45,14 @@ def get_option_setter(dataset_name):
 
 
 def create_dataset(opt):
-    """Create dataset given the option.
+    """Create a dataset given the option.
 
-    This function warps the class CustomDatasetDataLoader.
-    This is the main interface called by train.py and test.py.
+    This function wraps the class CustomDatasetDataLoader.
+        This is the main interface between this package and 'train.py'/'test.py'
 
     Example:
-    from data import create_dataset
-    dataset = create_dataset(opt)
+        >>> from data import create_dataset
+        >>> dataset = create_dataset(opt)
     """
     data_loader = CustomDatasetDataLoader(opt)
     dataset = data_loader.load_data()

data/base_dataset.py

@@ -11,18 +11,18 @@
 class BaseDataset(data.Dataset, ABC):
     """This class is an abstract base class (ABC) for datasets.
 
-    To create a subclass, you need to implement four functions:
-    -- <__init__> (initialize the class, first call BaseDataset.__init__(self, opt))
-    -- <__len__> (return the size of dataset)
-    -- <__getitem__>　(get a data point)
-    -- (optionally) <modify_commandline_options> (add dataset-specific options and set default options).
+    To create a subclass, you need to implement the following four functions:
+    -- <__init__>:                      initialize the class, first call BaseDataset.__init__(self, opt).
+    -- <__len__>:                       return the size of dataset.
+    -- <__getitem__>:                   get a data point.
+    -- <modify_commandline_options>:    (optionally) add dataset-specific options and set default options.
     """
 
     def __init__(self, opt):
         """Initialize the class; save the options in the class
 
         Parameters:
-            opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
+            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
         """
         self.opt = opt
         self.root = opt.dataroot
@@ -32,8 +32,8 @@ def modify_commandline_options(parser, is_train):
         """Add new dataset-specific options, and rewrite default values for existing options.
 
         Parameters:
-            parser -- original option parser
-            is_train -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
 
         Returns:
             the modified parser.

data/colorization_dataset.py

@@ -16,8 +16,15 @@ class ColorizationDataset(BaseDataset):
     def modify_commandline_options(parser, is_train):
         """Add new dataset-specific options, and rewrite default values for existing options.
 
+        Parameters:
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+
         By default, the number of channels for input image  is 1 (L) and
-         the nubmer of channels for output image is 2 (ab). The direction is from A to B
+        the nubmer of channels for output image is 2 (ab). The direction is from A to B
         """
         parser.set_defaults(input_nc=1, output_nc=2, direction='AtoB')
         return parser

data/single_dataset.py

@@ -13,7 +13,7 @@ def __init__(self, opt):
         """Initialize this dataset class.
 
         Parameters:
-            opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
         """
         BaseDataset.__init__(self, opt)
         self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))

data/template_dataset.py

@@ -23,8 +23,8 @@ def modify_commandline_options(parser, is_train):
         """Add new dataset-specific options, and rewrite default values for existing options.
 
         Parameters:
-            parser -- original option parser
-            is_train -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+            parser          -- original option parser
+            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
 
         Returns:
             the modified parser.
@@ -37,7 +37,7 @@ def __init__(self, opt):
         """Initialize this dataset class.
 
         Parameters:
-            opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
 
         A few things can be done here.
         - save the options (have been done in BaseDataset)

data/unaligned_dataset.py

@@ -20,7 +20,7 @@ def __init__(self, opt):
         """Initialize this dataset class.
 
         Parameters:
-            opt -- stores all the experiment flags; needs to be a subclass of BaseOptions
+            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
         """
         BaseDataset.__init__(self, opt)
         self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')  # create a path '/path/to/data/trainA'
@@ -40,13 +40,13 @@ def __getitem__(self, index):
         """Return a data point and its metadata information.
 
         Parameters:
-            index - - a random integer for data indexing
+            index (int)      -- a random integer for data indexing
 
         Returns a dictionary that contains A, B, A_paths and B_paths
-            A(tensor) - - an image in the input domain
-            B(tensor) - - its corresponding image in the target domain
-            A_paths(str) - - image paths
-            B_paths(str) - - image paths
+            A (tensor)       -- an image in the input domain
+            B (tensor)       -- its corresponding image in the target domain
+            A_paths (str)    -- image paths
+            B_paths (str)    -- image paths
         """
         A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
         if self.opt.serial_batches:   # make sure index is within then range

docs/overview.md

@@ -24,10 +24,10 @@ To help users better understand and use our codebase, we briefly overview the fu
 * [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 channel in [Lab](https://en.wikipedia.org/wiki/CIELAB_color_space) color space. By default, the model 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` discrimiator (PatchGAN  introduced by pix2pix), and a least-square GANs [objective](https://arxiv.org/abs/1611.04076) (`--gan_mode lsgan`).
+* [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 option 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.
+* [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,

models/__init__.py

@@ -1,17 +1,36 @@
+"""This package 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 the following five functions:
+    -- <__init__>:                      initialize the class; first call BaseModel.__init__(self, opt).
+    -- <set_input>:                     unpack data from dataset and apply preprocessing.
+    -- <forward>:                       produce intermediate results.
+    -- <optimize_parameters>:           calculate loss, gradients, and update network weights.
+    -- <modify_commandline_options>:    (optionally) add model-specific options and set default options.
+
+In the function <__init__>, you need to define four lists:
+    -- self.loss_names (str list):          specify the training losses that you want to plot and save.
+    -- self.model_names (str list):         specify the images that you want to display and save.
+    -- self.visual_names (str list):        define networks used in our training.
+    -- self.optimizers (optimzier list):    define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.
+
+Now you can use the model class by specifying flag '--model dummy'.
+See our template model class 'template_model.py' for an example.
+"""
+
 import importlib
 from models.base_model import BaseModel
 
 
 def find_model_using_name(model_name):
-    # Given the option --model [modelname],
-    # the file "models/modelname_model.py"
-    # will be imported.
+    """Import the module "models/[model_name]_model.py".
+
+    In the file, the class called DatasetNameModel() will
+    be instantiated. It has to be a subclass of BaseModel,
+    and it is case-insensitive.
+    """
     model_filename = "models." + model_name + "_model"
     modellib = importlib.import_module(model_filename)
-
-    # In the file, the class called ModelNameModel() will
-    # be instantiated. It has to be a subclass of BaseModel,
-    # and it is case-insensitive.
     model = None
     target_model_name = model_name.replace('_', '') + 'model'
     for name, cls in modellib.__dict__.items():
@@ -27,11 +46,21 @@ def find_model_using_name(model_name):
 
 
 def get_option_setter(model_name):
+    """Return the static method <modify_commandline_options> of the model class."""
     model_class = find_model_using_name(model_name)
     return model_class.modify_commandline_options
 
 
 def create_model(opt):
+    """Create a model given the option.
+
+    This function warps the class CustomDatasetDataLoader.
+    This is the main interface between this package and 'train.py'/'test.py'
+
+    Example:
+        >>> from models import create_model
+        >>> model = create_model(opt)
+    """
     model = find_model_using_name(opt.model)
     instance = model(opt)
     print("model [%s] was created" % type(instance).__name__)