Repository: keras-team/autokeras

Issue:
Docstring for model_trainer.py and preprocessor.py
### Feature Description
Clear and expressive document for model_trainer.py and preprocessor.py

### Reason
It's necessary to have documents for code to make the code more understandable

### Solution
Do it




Discussion:

Resolving patch:
diff --git a/autokeras/nn/model_trainer.py b/autokeras/nn/model_trainer.py
--- a/autokeras/nn/model_trainer.py
+++ b/autokeras/nn/model_trainer.py
@@ -13,6 +13,23 @@
 
 
 class ModelTrainerBase(abc.ABC):
+    """ A base class all model trainers will inherit from.
+
+    Attributes:
+        device: A string. Indicating the device to use. 'cuda' or 'cpu'.
+        train_loader: Training data wrapped in batches in Pytorch Dataloader.
+        test_loader: Testing data wrapped in batches in Pytorch Dataloader.
+        loss_function: A function with two parameters (prediction, target).
+            There is no specific requirement for the types of the parameters,
+            as long as they are compatible with the model and the data loaders.
+            The prediction should be the output of the model for a batch.
+            The target should be a batch of targets packed in the data loaders.
+        metric: It should be a subclass of class autokeras.metric.Metric.
+            In the compute(prediction, target) function, prediction and targets are,
+            all numpy arrays converted from the output of the model and the targets packed in the data loaders.
+        verbose: Verbosity mode.
+    """
+
     def __init__(self,
                  loss_function,
                  train_data,
@@ -30,6 +47,13 @@ def __init__(self,
     def train_model(self,
                     max_iter_num=Constant.MAX_ITER_NUM,
                     max_no_improvement_num=Constant.MAX_NO_IMPROVEMENT_NUM):
+        """Train the model.
+
+        Args:
+            max_iter_num: int, maximum numer of iteration
+            max_no_improvement_num: after max_no_improvement_num,
+                if the model still makes no improvement, finish training.
+        """
         pass
 
 
@@ -41,25 +65,14 @@ class ModelTrainer(ModelTrainerBase):
     Please see the details in the Attributes.
 
     Attributes:
-        device: A string. Indicating the device to use. 'cuda' or 'cpu'.
+        temp_model_path: Specify the path where temp model should be stored.
         model: An instance of Pytorch Module. The model that will be trained.
-        train_loader: Training data wrapped in batches in Pytorch Dataloader.
-        test_loader: Testing data wrapped in batches in Pytorch Dataloader.
-        loss_function: A function with two parameters (prediction, target).
-            There is no specific requirement for the types of the parameters,
-            as long as they are compatible with the model and the data loaders.
-            The prediction should be the output of the model for a batch.
-            The target should be a batch of targets packed in the data loaders.
-        optimizer: The optimizer is chosen to use the Pytorch Adam optimizer.
         early_stop: An instance of class EarlyStop.
-        metric: It should be a subclass of class autokeras.metric.Metric.
-            In the compute(prediction, target) function, prediction and targets are
-            all numpy arrays converted from the output of the model and the targets packed in the data loaders.
-        verbose: Verbosity mode.
+        optimizer: The optimizer is chosen to use the Pytorch Adam optimizer.
+        current_epoch: Record the current epoch.
     """
 
     def __init__(self, model, path, **kwargs):
-        """Init the ModelTrainer with `model`, `x_train`, `y_train`, `x_test`, `y_test`, `verbose`"""
         super().__init__(**kwargs)
         self.model = model
         self.model.to(self.device)
@@ -74,11 +87,16 @@ def train_model(self,
                     max_no_improvement_num=None):
         """Train the model.
 
+        Train the model with max_iter_num or max_no_improvement_num is met.
+
         Args:
             max_iter_num: An integer. The maximum number of epochs to train the model.
                 The training will stop when this number is reached.
             max_no_improvement_num: An integer. The maximum number of epochs when the loss value doesn't decrease.
                 The training will stop when this number is reached.
+
+        Returns:
+            A tuple of loss values and metric value.
         """
         if max_iter_num is None:
             max_iter_num = Constant.MAX_ITER_NUM
@@ -115,6 +133,7 @@ def train_model(self,
                 sum(test_metric_value_list[-last_num:]) / last_num)
 
     def _train(self):
+        """Where the actual train proceed."""
         self.model.train()
         loader = self.train_loader
         self.current_epoch += 1
@@ -147,6 +166,7 @@ def _train(self):
             progress_bar.close()
 
     def _test(self):
+        """Function for evaluation."""
         self.model.eval()
         test_loss = 0
         all_targets = []
@@ -173,6 +193,16 @@ def _load_model(self):
 
 
 class GANModelTrainer(ModelTrainerBase):
+    """A ModelTrainer especially for the GAN.
+
+    Attributes:
+        d_model: A discriminator model.
+        g_model: A generator model.
+        out_f: Out file.
+        out_size: Size of the output image.
+        optimizer_d: Optimizer for discriminator.
+        optimizer_g: Optimizer for generator.
+    """
     def __init__(self,
                  g_model,
                  d_model,
@@ -180,7 +210,16 @@ def __init__(self,
                  loss_function,
                  verbose,
                  gen_training_result=None):
-        """Init the ModelTrainer with `model`, `x_train`, `y_train`, `x_test`, `y_test`, `verbose`"""
+        """Initialize the GANModelTrainer.
+
+        Args:
+            g_model: The generator model to be trained.
+            d_model: The discriminator model to be trained.
+            train_data: the training data.
+            loss_function: The loss function for both discriminator and generator.
+            verbose: Whether to output the system output.
+            gen_training_result: Whether to generate the intermediate result while training.
+        """
         super().__init__(loss_function, train_data, verbose=verbose)
         self.d_model = d_model
         self.g_model = g_model
@@ -218,6 +257,7 @@ def train_model(self,
             progress_bar.close()
 
     def _train(self, epoch):
+        """Perform the actual train."""
         # put model into train mode
         self.d_model.train()
         # TODO: why?
@@ -276,6 +316,16 @@ def _train(self, epoch):
 
 
 class EarlyStop:
+    """A class check for early stop condition.
+
+    Attributes:
+        training_losses: Record all the training loss.
+        minimum_loss: The minimum loss we achieve so far. Used to compared to determine no improvement condition.
+        no_improvement_count: Current no improvement count.
+        _max_no_improvement_num: The maximum number specified.
+        _done: Whether condition met.
+        _min_loss_dec: A threshold for loss improvement.
+    """
     def __init__(self, max_no_improvement_num=Constant.MAX_NO_IMPROVEMENT_NUM, min_loss_dec=Constant.MIN_LOSS_DEC):
         super().__init__()
         self.training_losses = []
@@ -286,12 +336,26 @@ def __init__(self, max_no_improvement_num=Constant.MAX_NO_IMPROVEMENT_NUM, min_l
         self._min_loss_dec = min_loss_dec
 
     def on_train_begin(self):
+        """Initiate the early stop condition.
+
+        Call on every time the training iteration begins.
+        """
         self.training_losses = []
         self.no_improvement_count = 0
         self._done = False
         self.minimum_loss = float('inf')
 
     def on_epoch_end(self, loss):
+        """Check the early stop condition.
+
+        Call on every time the training iteration end.
+
+        Args:
+            loss: The loss function achieved by the epoch.
+
+        Returns:
+            True if condition met, otherwise False.
+        """
         self.training_losses.append(loss)
         if self._done and loss > (self.minimum_loss - self._min_loss_dec):
             return False
diff --git a/autokeras/preprocessor.py b/autokeras/preprocessor.py
--- a/autokeras/preprocessor.py
+++ b/autokeras/preprocessor.py
@@ -1,9 +1,11 @@
 import os
+from abc import ABC, abstractmethod
+
 import numpy as np
 import torch
 from torch.utils.data import Dataset, DataLoader
 from torchvision.transforms import ToPILImage, RandomCrop, RandomHorizontalFlip, ToTensor, Normalize, Compose
-from abc import ABC, abstractmethod
+
 from autokeras.constant import Constant
 from autokeras.utils import read_csv_file, read_image
 
@@ -54,6 +56,7 @@ def inverse_transform(self, data):
 
 class Cutout(object):
     """Randomly mask out one or more patches from an image.
+
     Args:
         n_holes (int): Number of patches to cut out of each image.
         length (int): The length (in pixels) of each square patch.
@@ -64,9 +67,11 @@ def __init__(self, n_holes, length):
         self.length = length
 
     def __call__(self, img):
-        """
+        """Perform the actual transformation.
+
         Args:
             img (Tensor): Tensor image of size (C, H, W).
+
         Returns:
             Tensor: Image with n_holes of dimension length x length cut out of it.
         """
@@ -94,18 +99,46 @@ def __call__(self, img):
 
 
 class DataTransformer(ABC):
+    """A superclass for all the DataTransformer."""
+    def __init__(self):
+        pass
+
     @abstractmethod
     def transform_train(self, data, targets=None, batch_size=None):
+        """ Transform the training data and get the DataLoader class.
+
+        Args:
+            data: x.
+            targets: y.
+            batch_size: the batchsize.
+
+        Returns:
+            dataloader: A torch.DataLoader class to represent the transformed data.
+        """
         raise NotImplementedError
 
     @abstractmethod
     def transform_test(self, data, target=None, batch_size=None):
+        """ Transform the training data and get the DataLoader class.
+
+        Args:
+            data: x.
+            targets: y.
+            batch_size: the batchsize.
+
+        Returns:
+            dataloader: A torch.DataLoader class to represent the transformed data.
+        """
         raise NotImplementedError
 
 
 class TextDataTransformer(DataTransformer):
+    """ A DataTransformer class for the text data."""
+    def __init__(self):
+        super().__init__()
 
     def transform_train(self, data, targets=None, batch_size=None):
+        """Transform the training dataset."""
         dataset = self._transform(compose_list=[], data=data, targets=targets)
         if batch_size is None:
             batch_size = Constant.MAX_BATCH_SIZE
@@ -114,6 +147,7 @@ def transform_train(self, data, targets=None, batch_size=None):
         return DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
     def transform_test(self, data, targets=None, batch_size=None):
+        """Transform the testing dataset."""
         dataset = self._transform(compose_list=[], data=data, targets=targets)
         if batch_size is None:
             batch_size = Constant.MAX_BATCH_SIZE
@@ -128,6 +162,15 @@ def _transform(self, compose_list, data, targets):
 
 
 class ImageDataTransformer(DataTransformer):
+    """ Perform basic image transformation and augmentation.
+
+    Attributes:
+        max_val: the maximum value of all data.
+        mean: the mean value.
+        std: the standard deviation.
+        augment: whether to perform augmentation on data.
+    """
+
     def __init__(self, data, augment=Constant.DATA_AUGMENTATION):
         self.max_val = data.max()
         data = data / self.max_val
@@ -136,6 +179,15 @@ def __init__(self, data, augment=Constant.DATA_AUGMENTATION):
         self.augment = augment
 
     def transform_train(self, data, targets=None, batch_size=None):
+        """ Transform the training data, perform random cropping data augmentation and basic random flip augmentation.
+
+        Args:
+            batch_size: int batch_size.
+            targets: the target of training set.
+
+        Returns:
+            A DataLoader class instance.
+        """
         short_edge_length = min(data.shape[1], data.shape[2])
         common_list = [Normalize(torch.Tensor(self.mean), torch.Tensor(self.std))]
         if self.augment:
@@ -160,6 +212,14 @@ def transform_train(self, data, targets=None, batch_size=None):
         return DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
     def transform_test(self, data, targets=None, batch_size=None):
+        """ Transform the test data, perform normalization.
+
+        Args:
+            batch_size: int batch_size.
+            targets: the target of test set.
+        Returns:
+            A DataLoader instance.
+        """
         common_list = [Normalize(torch.Tensor(self.mean), torch.Tensor(self.std))]
         compose_list = common_list
         if len(data.shape) != 4:
@@ -174,6 +234,16 @@ def transform_test(self, data, targets=None, batch_size=None):
         return DataLoader(dataset, batch_size=batch_size, shuffle=False)
 
     def _transform(self, compose_list, data, targets):
+        """Perform the actual transformation.
+
+        Args:
+            compose_list: a list of transforming operation.
+            data: x.
+            targets: y.
+
+        Returns:
+            A MultiTransformDataset class to represent the dataset.
+        """
         data = data / self.max_val
         args = [0, len(data.shape) - 1] + list(range(1, len(data.shape) - 1))
         data = torch.Tensor(data.transpose(*args))
@@ -182,6 +252,8 @@ def _transform(self, compose_list, data, targets):
 
 
 class MultiTransformDataset(Dataset):
+    """A class incorporate all transform method into a torch.Dataset class."""
+
     def __init__(self, dataset, target, compose):
         self.dataset = dataset
         self.target = target
@@ -198,6 +270,7 @@ def __len__(self):
 
 
 class BatchDataset(Dataset):
+    """A torch.Dataset class that can read data batch by batch."""
     def __init__(self, csv_file_path, image_path, has_target=True):
         file_names, target = read_csv_file(csv_file_path)
 


Test patch: