Merge pull request #2020 from microsoft/dev-refactor-doc

merge back to master: update doc index

src/sdk/pynni/nni/nas/pytorch/cdarts/mutator.py

@@ -127,18 +127,15 @@ class RegularizedMutatorParallel(DistributedDataParallel):
 class DartsDiscreteMutator(Mutator):
     """
     A mutator that applies the final sampling result of a parent mutator on another model to train.
+
+    Parameters
+    ----------
+    model : nn.Module
+        The model to apply the mutator.
+    parent_mutator : Mutator
+        The mutator that provides ``sample_final`` method, that will be called to get the architecture.
     """
     def __init__(self, model, parent_mutator):
-        """
-        Initialization.
-
-        Parameters
-        ----------
-        model : nn.Module
-            The model to apply the mutator.
-        parent_mutator : Mutator
-            The mutator that provides ``sample_final`` method, that will be called to get the architecture.
-        """
         super().__init__(model)
         self.__dict__["parent_mutator"] = parent_mutator  # avoid parameters to be included
 

src/sdk/pynni/nni/nas/pytorch/cdarts/trainer.py

@@ -32,73 +32,73 @@ class InteractiveKLLoss(nn.Module):
 
 
 class CdartsTrainer(object):
+    """
+    CDARTS trainer.
+
+    Parameters
+    ----------
+    model_small : nn.Module
+        PyTorch model to be trained. This is the search network of CDARTS.
+    model_large : nn.Module
+        PyTorch model to be trained. This is the evaluation network of CDARTS.
+    criterion : callable
+        Receives logits and ground truth label, return a loss tensor, e.g., ``nn.CrossEntropyLoss()``.
+    loaders : list of torch.utils.data.DataLoader
+        List of train data and valid data loaders, for training weights and architecture weights respectively.
+    samplers : list of torch.utils.data.Sampler
+        List of train data and valid data samplers. This can be PyTorch standard samplers if not distributed.
+        In distributed mode, sampler needs to have ``set_epoch`` method. Refer to data utils in CDARTS example for details.
+    logger : logging.Logger
+        The logger for logging. Will use nni logger by default (if logger is ``None``).
+    regular_coeff : float
+        The coefficient of regular loss.
+    regular_ratio : float
+        The ratio of regular loss.
+    warmup_epochs : int
+        The epochs to warmup the search network
+    fix_head : bool
+        ``True`` if fixing the paramters of auxiliary heads, else unfix the paramters of auxiliary heads.
+    epochs : int
+        Number of epochs planned for training.
+    steps_per_epoch : int
+        Steps of one epoch.
+    loss_alpha : float
+        The loss coefficient.
+    loss_T : float
+        The loss coefficient.
+    distributed : bool
+        ``True`` if using distributed training, else non-distributed training.
+    log_frequency : int
+        Step count per logging.
+    grad_clip : float
+        Gradient clipping for weights.
+    interactive_type : string
+        ``kl`` or ``smoothl1``.
+    output_path : string
+        Log storage path.
+    w_lr : float
+        Learning rate of the search network parameters.
+    w_momentum : float
+        Momentum of the search and the evaluation network.
+    w_weight_decay : float
+        The weight decay the search and the evaluation network parameters.
+    alpha_lr : float
+        Learning rate of the architecture parameters.
+    alpha_weight_decay : float
+        The weight decay the architecture parameters.
+    nasnet_lr : float
+        Learning rate of the evaluation network parameters.
+    local_rank : int
+        The number of thread.
+    share_module : bool
+        ``True`` if sharing the stem and auxiliary heads, else not sharing these modules.
+    """
     def __init__(self, model_small, model_large, criterion, loaders, samplers, logger=None,
                  regular_coeff=5, regular_ratio=0.2, warmup_epochs=2, fix_head=True,
                  epochs=32, steps_per_epoch=None, loss_alpha=2, loss_T=2, distributed=True,
                  log_frequency=10, grad_clip=5.0, interactive_type='kl', output_path='./outputs',
                  w_lr=0.2, w_momentum=0.9, w_weight_decay=3e-4, alpha_lr=0.2, alpha_weight_decay=1e-4,
                  nasnet_lr=0.2, local_rank=0, share_module=True):
-        """
-        Initialize a CdartsTrainer.
-
-        Parameters
-        ----------
-        model_small : nn.Module
-            PyTorch model to be trained. This is the search network of CDARTS.
-        model_large : nn.Module
-            PyTorch model to be trained. This is the evaluation network of CDARTS.
-        criterion : callable
-            Receives logits and ground truth label, return a loss tensor, e.g., ``nn.CrossEntropyLoss()``.
-        loaders : list of torch.utils.data.DataLoader
-            List of train data and valid data loaders, for training weights and architecture weights respectively.
-        samplers : list of torch.utils.data.Sampler
-            List of train data and valid data samplers. This can be PyTorch standard samplers if not distributed.
-            In distributed mode, sampler needs to have ``set_epoch`` method. Refer to data utils in CDARTS example for details.
-        logger : logging.Logger
-            The logger for logging. Will use nni logger by default (if logger is ``None``).
-        regular_coeff : float
-            The coefficient of regular loss.
-        regular_ratio : float
-            The ratio of regular loss.
-        warmup_epochs : int
-            The epochs to warmup the search network
-        fix_head : bool
-            ``True`` if fixing the paramters of auxiliary heads, else unfix the paramters of auxiliary heads.
-        epochs : int
-            Number of epochs planned for training.
-        steps_per_epoch : int
-            Steps of one epoch.
-        loss_alpha : float
-            The loss coefficient.
-        loss_T : float
-            The loss coefficient.
-        distributed : bool
-            ``True`` if using distributed training, else non-distributed training.
-        log_frequency : int
-            Step count per logging.
-        grad_clip : float
-            Gradient clipping for weights.
-        interactive_type : string
-            ``kl`` or ``smoothl1``.
-        output_path : string
-            Log storage path.
-        w_lr : float
-            Learning rate of the search network parameters.
-        w_momentum : float
-            Momentum of the search and the evaluation network.
-        w_weight_decay : float
-            The weight decay the search and the evaluation network parameters.
-        alpha_lr : float
-            Learning rate of the architecture parameters.
-        alpha_weight_decay : float
-            The weight decay the architecture parameters.
-        nasnet_lr : float
-            Learning rate of the evaluation network parameters.
-        local_rank : int
-            The number of thread.
-        share_module : bool
-            ``True`` if sharing the stem and auxiliary heads, else not sharing these modules.
-        """
         if logger is None:
             logger = logging.getLogger(__name__)
         train_loader, valid_loader = loaders

src/sdk/pynni/nni/nas/pytorch/classic_nas/mutator.py

@@ -22,12 +22,21 @@ INPUT_CHOICE = "input_choice"
 
 def get_and_apply_next_architecture(model):
     """
-    Wrapper of ClassicMutator to make it more meaningful,
-    similar to ```get_next_parameter``` for HPO.
+    Wrapper of :class:`~nni.nas.pytorch.classic_nas.mutator.ClassicMutator` to make it more meaningful,
+    similar to ``get_next_parameter`` for HPO.
+
+    Tt will generate search space based on ``model``.
+    If env ``NNI_GEN_SEARCH_SPACE`` exists, this is in dry run mode for
+    generating search space for the experiment.
+    If not, there are still two mode, one is nni experiment mode where users
+    use ``nnictl`` to start an experiment. The other is standalone mode
+    where users directly run the trial command, this mode chooses the first
+    one(s) for each LayerChoice and InputChoice.
+
     Parameters
     ----------
-    model : pytorch model
-        user's model with search space (e.g., LayerChoice, InputChoice) embedded in it
+    model : nn.Module
+        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
     """
     ClassicMutator(model)
 
@@ -36,23 +45,15 @@ class ClassicMutator(Mutator):
     """
     This mutator is to apply the architecture chosen from tuner.
     It implements the forward function of LayerChoice and InputChoice,
-    to only activate the chosen ones
+    to only activate the chosen ones.
+
+    Parameters
+    ----------
+    model : nn.Module
+        User's model with search space (e.g., LayerChoice, InputChoice) embedded in it.
     """
 
     def __init__(self, model):
-        """
-        Generate search space based on ```model```.
-        If env ```NNI_GEN_SEARCH_SPACE``` exists, this is in dry run mode for
-        generating search space for the experiment.
-        If not, there are still two mode, one is nni experiment mode where users
-        use ```nnictl``` to start an experiment. The other is standalone mode
-        where users directly run the trial command, this mode chooses the first
-        one(s) for each LayerChoice and InputChoice.
-        Parameters
-        ----------
-        model : PyTorch model
-            user's model with search space (e.g., LayerChoice, InputChoice) embedded in it
-        """
         super(ClassicMutator, self).__init__(model)
         self._chosen_arch = {}
         self._search_space = self._generate_search_space()
@@ -121,9 +122,15 @@ class ClassicMutator(Mutator):
         return torch.tensor(multihot_list, dtype=torch.bool)  # pylint: disable=not-callable
 
     def sample_search(self):
+        """
+        See :meth:`sample_final`.
+        """
         return self.sample_final()
 
     def sample_final(self):
+        """
+        Convert the chosen arch and apply it on model.
+        """
         assert set(self._chosen_arch.keys()) == set(self._search_space.keys()), \
             "Unmatched keys, expected keys '{}' from search space, found '{}'.".format(self._search_space.keys(),
                                                                                        self._chosen_arch.keys())

src/sdk/pynni/nni/nas/pytorch/darts/trainer.py

@@ -15,46 +15,46 @@ logger = logging.getLogger(__name__)
 
 
 class DartsTrainer(Trainer):
+    """
+    DARTS trainer.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model to be trained.
+    loss : callable
+        Receives logits and ground truth label, return a loss tensor.
+    metrics : callable
+        Receives logits and ground truth label, return a dict of metrics.
+    optimizer : Optimizer
+        The optimizer used for optimizing the model.
+    num_epochs : int
+        Number of epochs planned for training.
+    dataset_train : Dataset
+        Dataset for training. Will be split for training weights and architecture weights.
+    dataset_valid : Dataset
+        Dataset for testing.
+    mutator : DartsMutator
+        Use in case of customizing your own DartsMutator. By default will instantiate a DartsMutator.
+    batch_size : int
+        Batch size.
+    workers : int
+        Workers for data loading.
+    device : torch.device
+        ``torch.device("cpu")`` or ``torch.device("cuda")``.
+    log_frequency : int
+        Step count per logging.
+    callbacks : list of Callback
+        list of callbacks to trigger at events.
+    arc_learning_rate : float
+        Learning rate of architecture parameters.
+    unrolled : float
+        ``True`` if using second order optimization, else first order optimization.
+    """
     def __init__(self, model, loss, metrics,
                  optimizer, num_epochs, dataset_train, dataset_valid,
                  mutator=None, batch_size=64, workers=4, device=None, log_frequency=None,
                  callbacks=None, arc_learning_rate=3.0E-4, unrolled=False):
-        """
-        Initialize a DartsTrainer.
-
-        Parameters
-        ----------
-        model : nn.Module
-            PyTorch model to be trained.
-        loss : callable
-            Receives logits and ground truth label, return a loss tensor.
-        metrics : callable
-            Receives logits and ground truth label, return a dict of metrics.
-        optimizer : Optimizer
-            The optimizer used for optimizing the model.
-        num_epochs : int
-            Number of epochs planned for training.
-        dataset_train : Dataset
-            Dataset for training. Will be split for training weights and architecture weights.
-        dataset_valid : Dataset
-            Dataset for testing.
-        mutator : DartsMutator
-            Use in case of customizing your own DartsMutator. By default will instantiate a DartsMutator.
-        batch_size : int
-            Batch size.
-        workers : int
-            Workers for data loading.
-        device : torch.device
-            ``torch.device("cpu")`` or ``torch.device("cuda")``.
-        log_frequency : int
-            Step count per logging.
-        callbacks : list of Callback
-            list of callbacks to trigger at events.
-        arc_learning_rate : float
-            Learning rate of architecture parameters.
-        unrolled : float
-            ``True`` if using second order optimization, else first order optimization.
-        """
         super().__init__(model, mutator if mutator is not None else DartsMutator(model),
                          loss, metrics, optimizer, num_epochs, dataset_train, dataset_valid,
                          batch_size, workers, device, log_frequency, callbacks)

src/sdk/pynni/nni/nas/pytorch/enas/mutator.py

@@ -28,38 +28,38 @@ class StackedLSTMCell(nn.Module):
 
 
 class EnasMutator(Mutator):
+    """
+    A mutator that mutates the graph with RL.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model.
+    lstm_size : int
+        Controller LSTM hidden units.
+    lstm_num_layers : int
+        Number of layers for stacked LSTM.
+    tanh_constant : float
+        Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``.
+    cell_exit_extra_step : bool
+        If true, RL controller will perform an extra step at the exit of each MutableScope, dump the hidden state
+        and mark it as the hidden state of this MutableScope. This is to align with the original implementation of paper.
+    skip_target : float
+        Target probability that skipconnect will appear.
+    temperature : float
+        Temperature constant that divides the logits.
+    branch_bias : float
+        Manual bias applied to make some operations more likely to be chosen.
+        Currently this is implemented with a hardcoded match rule that aligns with original repo.
+        If a mutable has a ``reduce`` in its key, all its op choices
+        that contains `conv` in their typename will receive a bias of ``+self.branch_bias`` initially; while others
+        receive a bias of ``-self.branch_bias``.
+    entropy_reduction : str
+        Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced.
+    """
 
     def __init__(self, model, lstm_size=64, lstm_num_layers=1, tanh_constant=1.5, cell_exit_extra_step=False,
                  skip_target=0.4, temperature=None, branch_bias=0.25, entropy_reduction="sum"):
-        """
-        Initialize a EnasMutator.
-
-        Parameters
-        ----------
-        model : nn.Module
-            PyTorch model.
-        lstm_size : int
-            Controller LSTM hidden units.
-        lstm_num_layers : int
-            Number of layers for stacked LSTM.
-        tanh_constant : float
-            Logits will be equal to ``tanh_constant * tanh(logits)``. Don't use ``tanh`` if this value is ``None``.
-        cell_exit_extra_step : bool
-            If true, RL controller will perform an extra step at the exit of each MutableScope, dump the hidden state
-            and mark it as the hidden state of this MutableScope. This is to align with the original implementation of paper.
-        skip_target : float
-            Target probability that skipconnect will appear.
-        temperature : float
-            Temperature constant that divides the logits.
-        branch_bias : float
-            Manual bias applied to make some operations more likely to be chosen.
-            Currently this is implemented with a hardcoded match rule that aligns with original repo.
-            If a mutable has a ``reduce`` in its key, all its op choices
-            that contains `conv` in their typename will receive a bias of ``+self.branch_bias`` initially; while others
-            receive a bias of ``-self.branch_bias``.
-        entropy_reduction : str
-            Can be one of ``sum`` and ``mean``. How the entropy of multi-input-choice is reduced.
-        """
         super().__init__(model)
         self.lstm_size = lstm_size
         self.lstm_num_layers = lstm_num_layers

src/sdk/pynni/nni/nas/pytorch/enas/trainer.py

@@ -16,64 +16,64 @@ logger = logging.getLogger(__name__)
 
 
 class EnasTrainer(Trainer):
+    """
+    ENAS trainer.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model to be trained.
+    loss : callable
+        Receives logits and ground truth label, return a loss tensor.
+    metrics : callable
+        Receives logits and ground truth label, return a dict of metrics.
+    reward_function : callable
+        Receives logits and ground truth label, return a tensor, which will be feeded to RL controller as reward.
+    optimizer : Optimizer
+        The optimizer used for optimizing the model.
+    num_epochs : int
+        Number of epochs planned for training.
+    dataset_train : Dataset
+        Dataset for training. Will be split for training weights and architecture weights.
+    dataset_valid : Dataset
+        Dataset for testing.
+    mutator : EnasMutator
+        Use when customizing your own mutator or a mutator with customized parameters.
+    batch_size : int
+        Batch size.
+    workers : int
+        Workers for data loading.
+    device : torch.device
+        ``torch.device("cpu")`` or ``torch.device("cuda")``.
+    log_frequency : int
+        Step count per logging.
+    callbacks : list of Callback
+        list of callbacks to trigger at events.
+    entropy_weight : float
+        Weight of sample entropy loss.
+    skip_weight : float
+        Weight of skip penalty loss.
+    baseline_decay : float
+        Decay factor of baseline. New baseline will be equal to ``baseline_decay * baseline_old + reward * (1 - baseline_decay)``.
+    child_steps : int
+        How many mini-batches for model training per epoch.
+    mutator_lr : float
+        Learning rate for RL controller.
+    mutator_steps_aggregate : int
+        Number of steps that will be aggregated into one mini-batch for RL controller.
+    mutator_steps : int
+        Number of mini-batches for each epoch of RL controller learning.
+    aux_weight : float
+        Weight of auxiliary head loss. ``aux_weight * aux_loss`` will be added to total loss.
+    test_arc_per_epoch : int
+        How many architectures are chosen for direct test after each epoch.
+    """
     def __init__(self, model, loss, metrics, reward_function,
                  optimizer, num_epochs, dataset_train, dataset_valid,
                  mutator=None, batch_size=64, workers=4, device=None, log_frequency=None, callbacks=None,
                  entropy_weight=0.0001, skip_weight=0.8, baseline_decay=0.999, child_steps=500,
                  mutator_lr=0.00035, mutator_steps_aggregate=20, mutator_steps=50, aux_weight=0.4,
                  test_arc_per_epoch=1):
-        """
-        Initialize an EnasTrainer.
-
-        Parameters
-        ----------
-        model : nn.Module
-            PyTorch model to be trained.
-        loss : callable
-            Receives logits and ground truth label, return a loss tensor.
-        metrics : callable
-            Receives logits and ground truth label, return a dict of metrics.
-        reward_function : callable
-            Receives logits and ground truth label, return a tensor, which will be feeded to RL controller as reward.
-        optimizer : Optimizer
-            The optimizer used for optimizing the model.
-        num_epochs : int
-            Number of epochs planned for training.
-        dataset_train : Dataset
-            Dataset for training. Will be split for training weights and architecture weights.
-        dataset_valid : Dataset
-            Dataset for testing.
-        mutator : EnasMutator
-            Use when customizing your own mutator or a mutator with customized parameters.
-        batch_size : int
-            Batch size.
-        workers : int
-            Workers for data loading.
-        device : torch.device
-            ``torch.device("cpu")`` or ``torch.device("cuda")``.
-        log_frequency : int
-            Step count per logging.
-        callbacks : list of Callback
-            list of callbacks to trigger at events.
-        entropy_weight : float
-            Weight of sample entropy loss.
-        skip_weight : float
-            Weight of skip penalty loss.
-        baseline_decay : float
-            Decay factor of baseline. New baseline will be equal to ``baseline_decay * baseline_old + reward * (1 - baseline_decay)``.
-        child_steps : int
-            How many mini-batches for model training per epoch.
-        mutator_lr : float
-            Learning rate for RL controller.
-        mutator_steps_aggregate : int
-            Number of steps that will be aggregated into one mini-batch for RL controller.
-        mutator_steps : int
-            Number of mini-batches for each epoch of RL controller learning.
-        aux_weight : float
-            Weight of auxiliary head loss. ``aux_weight * aux_loss`` will be added to total loss.
-        test_arc_per_epoch : int
-            How many architectures are chosen for direct test after each epoch.
-        """
         super().__init__(model, mutator if mutator is not None else EnasMutator(model),
                          loss, metrics, optimizer, num_epochs, dataset_train, dataset_valid,
                          batch_size, workers, device, log_frequency, callbacks)

src/sdk/pynni/nni/nas/pytorch/fixed.py

@@ -10,20 +10,20 @@ from nni.nas.pytorch.mutator import Mutator
 
 
 class FixedArchitecture(Mutator):
+    """
+    Fixed architecture mutator that always selects a certain graph.
+
+    Parameters
+    ----------
+    model : nn.Module
+        A mutable network.
+    fixed_arc : str or dict
+        Path to the architecture checkpoint (a string), or preloaded architecture object (a dict).
+    strict : bool
+        Force everything that appears in ``fixed_arc`` to be used at least once.
+    """
 
     def __init__(self, model, fixed_arc, strict=True):
-        """
-        Initialize a fixed architecture mutator.
-
-        Parameters
-        ----------
-        model : nn.Module
-            A mutable network.
-        fixed_arc : str or dict
-            Path to the architecture checkpoint (a string), or preloaded architecture object (a dict).
-        strict : bool
-            Force everything that appears in `fixed_arc` to be used at least once.
-        """
         super().__init__(model)
         self._fixed_arc = fixed_arc
 
@@ -35,9 +35,15 @@ class FixedArchitecture(Mutator):
             raise RuntimeError("Missing keys in fixed architecture: {}.".format(mutable_keys - fixed_arc_keys))
 
     def sample_search(self):
+        """
+        Always returns the fixed architecture.
+        """
         return self._fixed_arc
 
     def sample_final(self):
+        """
+        Always returns the fixed architecture.
+        """
         return self._fixed_arc
 
 
@@ -66,6 +72,7 @@ def apply_fixed_architecture(model, fixed_arc):
     Returns
     -------
     FixedArchitecture
+        Mutator that is responsible for fixes the graph.
     """
 
     if isinstance(fixed_arc, str):

src/sdk/pynni/nni/nas/pytorch/mutables.py

@@ -19,7 +19,19 @@ class Mutable(nn.Module):
     decisions among different mutables. In mutator's implementation, mutators should use the key to
     distinguish different mutables. Mutables that share the same key should be "similar" to each other.
 
-    Currently the default scope for keys is global.
+    Currently the default scope for keys is global. By default, the keys uses a global counter from 1 to
+    produce unique ids.
+
+    Parameters
+    ----------
+    key : str
+        The key of mutable.
+
+    Notes
+    -----
+    The counter is program level, but mutables are model level. In case multiple models are defined, and
+    you want to have `counter` starting from 1 in the second model, it's recommended to assign keys manually
+    instead of using automatic keys.
     """
 
     def __init__(self, key=None):
@@ -51,10 +63,16 @@ class Mutable(nn.Module):
 
     @property
     def key(self):
+        """
+        Read-only property of key.
+        """
         return self._key
 
     @property
     def name(self):
+        """
+        After the search space is parsed, it will be the module name of the mutable.
+        """
         return self._name if hasattr(self, "_name") else "_key"
 
     @name.setter
@@ -75,11 +93,23 @@ class Mutable(nn.Module):
 class MutableScope(Mutable):
     """
     Mutable scope marks a subgraph/submodule to help mutators make better decisions.
-    Mutators get notified when a mutable scope is entered and exited. Mutators can override ``enter_mutable_scope``
-    and ``exit_mutable_scope`` to catch corresponding events, and do status dump or update.
-    MutableScope are also mutables that are listed in the mutables (search space).
-    """
 
+    If not annotated with mutable scope, search space will be flattened as a list. However, some mutators might
+    need to leverage the concept of a "cell". So if a module is defined as a mutable scope, everything in it will
+    look like "sub-search-space" in the scope. Scopes can be nested.
+
+    There are two ways mutators can use mutable scope. One is to traverse the search space as a tree during initialization
+    and reset. The other is to implement `enter_mutable_scope` and `exit_mutable_scope`. They are called before and after
+    the forward method of the class inheriting mutable scope.
+
+    Mutable scopes are also mutables that are listed in the mutator.mutables (search space), but they are not supposed
+    to appear in the dict of choices.
+
+    Parameters
+    ----------
+    key : str
+        Key of mutable scope.
+    """
     def __init__(self, key):
         super().__init__(key=key)
 
@@ -93,6 +123,31 @@ class MutableScope(Mutable):
 
 
 class LayerChoice(Mutable):
+    """
+    Layer choice selects one of the ``op_candidates``, then apply it on inputs and return results.
+    In rare cases, it can also select zero or many.
+
+    Layer choice does not allow itself to be nested.
+
+    Parameters
+    ----------
+    op_candidates : list of nn.Module
+        A module list to be selected from.
+    reduction : str
+        ``mean``, ``concat``, ``sum`` or ``none``. Policy if multiples are selected.
+        If ``none``, a list is returned. ``mean`` returns the average. ``sum`` returns the sum.
+        ``concat`` concatenate the list at dimension 1.
+    return_mask : bool
+        If ``return_mask``, return output tensor and a mask. Otherwise return tensor only.
+    key : str
+        Key of the input choice.
+
+    Attributes
+    ----------
+    length : int
+        Number of ops to choose from.
+    """
+
     def __init__(self, op_candidates, reduction="sum", return_mask=False, key=None):
         super().__init__(key=key)
         self.length = len(op_candidates)
@@ -101,6 +156,12 @@ class LayerChoice(Mutable):
         self.return_mask = return_mask
 
     def forward(self, *inputs):
+        """
+        Returns
+        -------
+        tuple of tensors
+            Output and selection mask. If ``return_mask`` is ``False``, only output is returned.
+        """
         out, mask = self.mutator.on_forward_layer_choice(self, *inputs)
         if self.return_mask:
             return out, mask
@@ -109,42 +170,62 @@ class LayerChoice(Mutable):
 
 class InputChoice(Mutable):
     """
-    Input choice selects `n_chosen` inputs from `choose_from` (contains `n_candidates` keys). For beginners,
-    use `n_candidates` instead of `choose_from` is a safe option. To get the most power out of it, you might want to
-    know about `choose_from`.
+    Input choice selects ``n_chosen`` inputs from ``choose_from`` (contains ``n_candidates`` keys). For beginners,
+    use ``n_candidates`` instead of ``choose_from`` is a safe option. To get the most power out of it, you might want to
+    know about ``choose_from``.
 
-    The keys in `choose_from` can be keys that appear in past mutables, or ``NO_KEY`` if there are no suitable ones.
+    The keys in ``choose_from`` can be keys that appear in past mutables, or ``NO_KEY`` if there are no suitable ones.
     The keys are designed to be the keys of the sources. To help mutators make better decisions,
     mutators might be interested in how the tensors to choose from come into place. For example, the tensor is the
     output of some operator, some node, some cell, or some module. If this operator happens to be a mutable (e.g.,
     ``LayerChoice`` or ``InputChoice``), it has a key naturally that can be used as a source key. If it's a
-    module/submodule, it needs to be annotated with a key: that's where a ``MutableScope`` is needed.
+    module/submodule, it needs to be annotated with a key: that's where a :class:`MutableScope` is needed.
+
+    In the example below, ``input_choice`` is a 4-choose-any. The first 3 is semantically output of cell1, output of cell2,
+    output of cell3 with respectively. Notice that an extra max pooling is followed by cell1, indicating x1 is not
+    "actually" the direct output of cell1.
+
+    .. code-block:: python
+
+        class Cell(MutableScope):
+            pass
+
+        class Net(nn.Module):
+            def __init__(self):
+                self.cell1 = Cell("cell1")
+                self.cell2 = Cell("cell2")
+                self.op = LayerChoice([conv3x3(), conv5x5()], key="op")
+                self.input_choice = InputChoice(choose_from=["cell1", "cell2", "op", InputChoice.NO_KEY])
+
+            def forward(self, x):
+                x1 = max_pooling(self.cell1(x))
+                x2 = self.cell2(x)
+                x3 = self.op(x)
+                x4 = torch.zeros_like(x)
+                return self.input_choice([x1, x2, x3, x4])
+
+    Parameters
+    ----------
+    n_candidates : int
+        Number of inputs to choose from.
+    choose_from : list of str
+        List of source keys to choose from. At least of one of ``choose_from`` and ``n_candidates`` must be fulfilled.
+        If ``n_candidates`` has a value but ``choose_from`` is None, it will be automatically treated as ``n_candidates``
+        number of empty string.
+    n_chosen : int
+        Recommended inputs to choose. If None, mutator is instructed to select any.
+    reduction : str
+        ``mean``, ``concat``, ``sum`` or ``none``. See :class:`LayerChoice`.
+    return_mask : bool
+        If ``return_mask``, return output tensor and a mask. Otherwise return tensor only.
+    key : str
+        Key of the input choice.
     """
 
     NO_KEY = ""
 
     def __init__(self, n_candidates=None, choose_from=None, n_chosen=None,
                  reduction="sum", return_mask=False, key=None):
-        """
-        Initialization.
-
-        Parameters
-        ----------
-        n_candidates : int
-            Number of inputs to choose from.
-        choose_from : list of str
-            List of source keys to choose from. At least of one of `choose_from` and `n_candidates` must be fulfilled.
-            If `n_candidates` has a value but `choose_from` is None, it will be automatically treated as `n_candidates`
-            number of empty string.
-        n_chosen : int
-            Recommended inputs to choose. If None, mutator is instructed to select any.
-        reduction : str
-            `mean`, `concat`, `sum` or `none`.
-        return_mask : bool
-            If `return_mask`, return output tensor and a mask. Otherwise return tensor only.
-        key : str
-            Key of the input choice.
-        """
         super().__init__(key=key)
         # precondition check
         assert n_candidates is not None or choose_from is not None, "At least one of `n_candidates` and `choose_from`" \
@@ -172,12 +253,13 @@ class InputChoice(Mutable):
         ----------
         optional_inputs : list or dict
             Recommended to be a dict. As a dict, inputs will be converted to a list that follows the order of
-            `choose_from` in initialization. As a list, inputs must follow the semantic order that is the same as
-            `choose_from`.
+            ``choose_from`` in initialization. As a list, inputs must follow the semantic order that is the same as
+            ``choose_from``.
 
         Returns
         -------
-        tuple of torch.Tensor and torch.Tensor or torch.Tensor
+        tuple of tensors
+            Output and selection mask. If ``return_mask`` is ``False``, only output is returned.
         """
         optional_input_list = optional_inputs
         if isinstance(optional_inputs, dict):

src/sdk/pynni/nni/nas/pytorch/mutator.py

@@ -43,10 +43,6 @@ class Mutator(BaseMutator):
         """
         Reset the mutator by call the `sample_search` to resample (for search). Stores the result in a local
         variable so that `on_forward_layer_choice` and `on_forward_input_choice` can use the decision directly.
-
-        Returns
-        -------
-        None
         """
         self._cache = self.sample_search()
 
@@ -57,25 +53,28 @@ class Mutator(BaseMutator):
         Returns
         -------
         dict
+            A mapping from key of mutables to decisions.
         """
         return self.sample_final()
 
     def on_forward_layer_choice(self, mutable, *inputs):
         """
-        On default, this method calls :meth:`on_calc_layer_choice_mask` to get a mask on how to choose between layers
-        (either by switch or by weights), then it will reduce the list of all tensor outputs with the policy specified
-        in `mutable.reduction`. It will also cache the mask with corresponding `mutable.key`.
+        On default, this method retrieves the decision obtained previously, and select certain operations.
+        Only operations with non-zero weight will be executed. The results will be added to a list.
+        Then it will reduce the list of all tensor outputs with the policy specified in `mutable.reduction`.
 
         Parameters
         ----------
         mutable : LayerChoice
+            Layer choice module.
         inputs : list of torch.Tensor
+            Inputs
 
         Returns
         -------
         tuple of torch.Tensor and torch.Tensor
+            Output and mask.
         """
-
         def _map_fn(op, *inputs):
             return op(*inputs)
 
@@ -87,20 +86,20 @@ class Mutator(BaseMutator):
 
     def on_forward_input_choice(self, mutable, tensor_list):
         """
-        On default, this method calls :meth:`on_calc_input_choice_mask` with `tags`
-        to get a mask on how to choose between inputs (either by switch or by weights), then it will reduce
-        the list of all tensor outputs with the policy specified in `mutable.reduction`. It will also cache the
-        mask with corresponding `mutable.key`.
+        On default, this method retrieves the decision obtained previously, and select certain tensors.
+        Then it will reduce the list of all tensor outputs with the policy specified in `mutable.reduction`.
 
         Parameters
         ----------
         mutable : InputChoice
+            Input choice module.
         tensor_list : list of torch.Tensor
-        tags : list of string
+            Tensor list to apply the decision on.
 
         Returns
         -------
         tuple of torch.Tensor and torch.Tensor
+            Output and mask.
         """
         mask = self._get_decision(mutable)
         assert len(mask) == mutable.n_candidates, \

src/sdk/pynni/nni/nas/pytorch/random/mutator.py

@@ -6,7 +6,15 @@ from nni.nas.pytorch.mutables import LayerChoice, InputChoice
 
 
 class RandomMutator(Mutator):
+    """
+    Random mutator that samples a random candidate in the search space each time ``reset()``.
+    It uses random function in PyTorch, so users can set seed in PyTorch to ensure deterministic behavior.
+    """
+
     def sample_search(self):
+        """
+        Sample a random candidate.
+        """
         result = dict()
         for mutable in self.mutables:
             if isinstance(mutable, LayerChoice):
@@ -22,4 +30,7 @@ class RandomMutator(Mutator):
         return result
 
     def sample_final(self):
+        """
+        Same as :meth:`sample_search`.
+        """
         return self.sample_search()

src/sdk/pynni/nni/nas/pytorch/spos/evolution.py

@@ -16,29 +16,29 @@ _logger = logging.getLogger(__name__)
 
 
 class SPOSEvolution(Tuner):
+    """
+    SPOS evolution tuner.
+
+    Parameters
+    ----------
+    max_epochs : int
+        Maximum number of epochs to run.
+    num_select : int
+        Number of survival candidates of each epoch.
+    num_population : int
+        Number of candidates at the start of each epoch. If candidates generated by
+        crossover and mutation are not enough, the rest will be filled with random
+        candidates.
+    m_prob : float
+        The probability of mutation.
+    num_crossover : int
+        Number of candidates generated by crossover in each epoch.
+    num_mutation : int
+        Number of candidates generated by mutation in each epoch.
+    """
 
     def __init__(self, max_epochs=20, num_select=10, num_population=50, m_prob=0.1,
                  num_crossover=25, num_mutation=25):
-        """
-        Initialize SPOS Evolution Tuner.
-
-        Parameters
-        ----------
-        max_epochs : int
-            Maximum number of epochs to run.
-        num_select : int
-            Number of survival candidates of each epoch.
-        num_population : int
-            Number of candidates at the start of each epoch. If candidates generated by
-            crossover and mutation are not enough, the rest will be filled with random
-            candidates.
-        m_prob : float
-            The probability of mutation.
-        num_crossover : int
-            Number of candidates generated by crossover in each epoch.
-        num_mutation : int
-            Number of candidates generated by mutation in each epoch.
-        """
         assert num_population >= num_select
         self.max_epochs = max_epochs
         self.num_select = num_select

src/sdk/pynni/nni/nas/pytorch/spos/mutator.py

@@ -10,27 +10,29 @@ _logger = logging.getLogger(__name__)
 
 
 class SPOSSupernetTrainingMutator(RandomMutator):
+    """
+    A random mutator with flops limit.
+
+    Parameters
+    ----------
+    model : nn.Module
+        PyTorch model.
+    flops_func : callable
+        Callable that takes a candidate from `sample_search` and returns its candidate. When `flops_func`
+        is None, functions related to flops will be deactivated.
+    flops_lb : number
+        Lower bound of flops.
+    flops_ub : number
+        Upper bound of flops.
+    flops_bin_num : number
+        Number of bins divided for the interval of flops to ensure the uniformity. Bigger number will be more
+        uniform, but the sampling will be slower.
+    flops_sample_timeout : int
+        Maximum number of attempts to sample before giving up and use a random candidate.
+    """
     def __init__(self, model, flops_func=None, flops_lb=None, flops_ub=None,
                  flops_bin_num=7, flops_sample_timeout=500):
-        """
 
-        Parameters
-        ----------
-        model : nn.Module
-            PyTorch model.
-        flops_func : callable
-            Callable that takes a candidate from `sample_search` and returns its candidate. When `flops_func`
-            is None, functions related to flops will be deactivated.
-        flops_lb : number
-            Lower bound of flops.
-        flops_ub : number
-            Upper bound of flops.
-        flops_bin_num : number
-            Number of bins divided for the interval of flops to ensure the uniformity. Bigger number will be more
-            uniform, but the sampling will be slower.
-        flops_sample_timeout : int
-            Maximum number of attempts to sample before giving up and use a random candidate.
-        """
         super().__init__(model)
         self._flops_func = flops_func
         if self._flops_func is not None:

src/sdk/pynni/nni/nas/pytorch/spos/trainer.py

@@ -15,43 +15,42 @@ logger = logging.getLogger(__name__)
 class SPOSSupernetTrainer(Trainer):
     """
     This trainer trains a supernet that can be used for evolution search.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model with mutables.
+    mutator : Mutator
+        A mutator object that has been initialized with the model.
+    loss : callable
+        Called with logits and targets. Returns a loss tensor.
+    metrics : callable
+        Returns a dict that maps metrics keys to metrics data.
+    optimizer : Optimizer
+        Optimizer that optimizes the model.
+    num_epochs : int
+        Number of epochs of training.
+    train_loader : iterable
+        Data loader of training. Raise ``StopIteration`` when one epoch is exhausted.
+    dataset_valid : iterable
+        Data loader of validation. Raise ``StopIteration`` when one epoch is exhausted.
+    batch_size : int
+        Batch size.
+    workers: int
+        Number of threads for data preprocessing. Not used for this trainer. Maybe removed in future.
+    device : torch.device
+        Device object. Either ``torch.device("cuda")`` or ``torch.device("cpu")``. When ``None``, trainer will
+        automatic detects GPU and selects GPU first.
+    log_frequency : int
+        Number of mini-batches to log metrics.
+    callbacks : list of Callback
+        Callbacks to plug into the trainer. See Callbacks.
     """
 
     def __init__(self, model, loss, metrics,
                  optimizer, num_epochs, train_loader, valid_loader,
                  mutator=None, batch_size=64, workers=4, device=None, log_frequency=None,
                  callbacks=None):
-        """
-        Parameters
-        ----------
-        model : nn.Module
-            Model with mutables.
-        mutator : Mutator
-            A mutator object that has been initialized with the model.
-        loss : callable
-            Called with logits and targets. Returns a loss tensor.
-        metrics : callable
-            Returns a dict that maps metrics keys to metrics data.
-        optimizer : Optimizer
-            Optimizer that optimizes the model.
-        num_epochs : int
-            Number of epochs of training.
-        train_loader : iterable
-            Data loader of training. Raise ``StopIteration`` when one epoch is exhausted.
-        dataset_valid : iterable
-            Data loader of validation. Raise ``StopIteration`` when one epoch is exhausted.
-        batch_size : int
-            Batch size.
-        workers: int
-            Number of threads for data preprocessing. Not used for this trainer. Maybe removed in future.
-        device : torch.device
-            Device object. Either ``torch.device("cuda")`` or ``torch.device("cpu")``. When ``None``, trainer will
-            automatic detects GPU and selects GPU first.
-        log_frequency : int
-            Number of mini-batches to log metrics.
-        callbacks : list of Callback
-            Callbacks to plug into the trainer. See Callbacks.
-        """
         assert torch.cuda.is_available()
         super().__init__(model, mutator if mutator is not None else SPOSSupernetTrainingMutator(model),
                          loss, metrics, optimizer, num_epochs, None, None,

src/sdk/pynni/nni/nas/pytorch/trainer.py

@@ -24,42 +24,54 @@ class TorchTensorEncoder(json.JSONEncoder):
 
 
 class Trainer(BaseTrainer):
+    """
+    A trainer with some helper functions implemented. To implement a new trainer,
+    users need to implement :meth:`train_one_epoch`, :meth:`validate_one_epoch` and :meth:`checkpoint`.
+
+    Parameters
+    ----------
+    model : nn.Module
+        Model with mutables.
+    mutator : BaseMutator
+        A mutator object that has been initialized with the model.
+    loss : callable
+        Called with logits and targets. Returns a loss tensor.
+        See `PyTorch loss functions`_ for examples.
+    metrics : callable
+        Called with logits and targets. Returns a dict that maps metrics keys to metrics data. For example,
+
+        .. code-block:: python
+
+            def metrics_fn(output, target):
+                return {"acc1": accuracy(output, target, topk=1), "acc5": accuracy(output, target, topk=5)}
+
+    optimizer : Optimizer
+        Optimizer that optimizes the model.
+    num_epochs : int
+        Number of epochs of training.
+    dataset_train : torch.utils.data.Dataset
+        Dataset of training. If not otherwise specified, ``dataset_train`` and ``dataset_valid`` should be standard
+        PyTorch Dataset. See `torch.utils.data`_ for examples.
+    dataset_valid : torch.utils.data.Dataset
+        Dataset of validation/testing.
+    batch_size : int
+        Batch size.
+    workers : int
+        Number of workers used in data preprocessing.
+    device : torch.device
+        Device object. Either ``torch.device("cuda")`` or ``torch.device("cpu")``. When ``None``, trainer will
+        automatic detects GPU and selects GPU first.
+    log_frequency : int
+        Number of mini-batches to log metrics.
+    callbacks : list of Callback
+        Callbacks to plug into the trainer. See Callbacks.
+
+
+    .. _`PyTorch loss functions`: https://pytorch.org/docs/stable/nn.html#loss-functions
+    .. _`torch.utils.data`: https://pytorch.org/docs/stable/data.html
+    """
     def __init__(self, model, mutator, loss, metrics, optimizer, num_epochs,
                  dataset_train, dataset_valid, batch_size, workers, device, log_frequency, callbacks):
-        """
-        Trainer initialization.
-
-        Parameters
-        ----------
-        model : nn.Module
-            Model with mutables.
-        mutator : BaseMutator
-            A mutator object that has been initialized with the model.
-        loss : callable
-            Called with logits and targets. Returns a loss tensor.
-        metrics : callable
-            Returns a dict that maps metrics keys to metrics data.
-        optimizer : Optimizer
-            Optimizer that optimizes the model.
-        num_epochs : int
-            Number of epochs of training.
-        dataset_train : torch.utils.data.Dataset
-            Dataset of training.
-        dataset_valid : torch.utils.data.Dataset
-            Dataset of validation/testing.
-        batch_size : int
-            Batch size.
-        workers : int
-            Number of workers used in data preprocessing.
-        device : torch.device
-            Device object. Either ``torch.device("cuda")`` or ``torch.device("cpu")``. When ``None``, trainer will
-            automatic detects GPU and selects GPU first.
-        log_frequency : int
-            Number of mini-batches to log metrics.
-        callbacks : list of Callback
-            Callbacks to plug into the trainer. See Callbacks.
-        """
-
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else device
         self.model = model
         self.mutator = mutator
@@ -84,13 +96,38 @@ class Trainer(BaseTrainer):
 
     @abstractmethod
     def train_one_epoch(self, epoch):
+        """
+        Train one epoch.
+
+        Parameters
+        ----------
+        epoch : int
+            Epoch number starting from 0.
+        """
         pass
 
     @abstractmethod
     def validate_one_epoch(self, epoch):
+        """
+        Validate one epoch.
+
+        Parameters
+        ----------
+        epoch : int
+            Epoch number starting from 0.
+        """
         pass
 
     def train(self, validate=True):
+        """
+        Train ``num_epochs``.
+        Trigger callbacks at the start and the end of each epoch.
+
+        Parameters
+        ----------
+        validate : bool
+            If ``true``, will do validation every epoch.
+        """
         for epoch in range(self.num_epochs):
             for callback in self.callbacks:
                 callback.on_epoch_begin(epoch)
@@ -108,12 +145,26 @@ class Trainer(BaseTrainer):
                 callback.on_epoch_end(epoch)
 
     def validate(self):
+        """
+        Do one validation.
+        """
         self.validate_one_epoch(-1)
 
     def export(self, file):
+        """
+        Call ``mutator.export()`` and dump the architecture to ``file``.
+
+        Parameters
+        ----------
+        file : str
+            A file path. Expected to be a JSON.
+        """
         mutator_export = self.mutator.export()
         with open(file, "w") as f:
             json.dump(mutator_export, f, indent=2, sort_keys=True, cls=TorchTensorEncoder)
 
     def checkpoint(self):
+        """
+        Return trainer checkpoint.
+        """
         raise NotImplementedError("Not implemented yet")

src/sdk/pynni/nni/nas/pytorch/utils.py

@@ -12,6 +12,9 @@ _logger = logging.getLogger(__name__)
 
 
 def global_mutable_counting():
+    """
+    A program level counter starting from 1.
+    """
     global _counter
     _counter += 1
     return _counter
@@ -26,6 +29,9 @@ def _reset_global_mutable_counting():
 
 
 def to_device(obj, device):
+    """
+    Move a tensor, tuple, list, or dict onto device.
+    """
     if torch.is_tensor(obj):
         return obj.to(device)
     if isinstance(obj, tuple):
@@ -40,12 +46,18 @@ def to_device(obj, device):
 
 
 class AverageMeterGroup:
-    """Average meter group for multiple average meters"""
+    """
+    Average meter group for multiple average meters.
+    """
 
     def __init__(self):
         self.meters = OrderedDict()
 
     def update(self, data):
+        """
+        Update the meter group with a dict of metrics.
+        Non-exist average meters will be automatically created.
+        """
         for k, v in data.items():
             if k not in self.meters:
                 self.meters[k] = AverageMeter(k, ":4f")
@@ -61,34 +73,49 @@ class AverageMeterGroup:
         return "  ".join(str(v) for v in self.meters.values())
 
     def summary(self):
+        """
+        Return a summary string of group data.
+        """
         return "  ".join(v.summary() for v in self.meters.values())
 
 
 class AverageMeter:
-    """Computes and stores the average and current value"""
+    """
+    Computes and stores the average and current value.
+
+    Parameters
+    ----------
+    name : str
+        Name to display.
+    fmt : str
+        Format string to print the values.
+    """
 
     def __init__(self, name, fmt=':f'):
-        """
-        Initialization of AverageMeter
-
-        Parameters
-        ----------
-        name : str
-            Name to display.
-        fmt : str
-            Format string to print the values.
-        """
         self.name = name
         self.fmt = fmt
         self.reset()
 
     def reset(self):
+        """
+        Reset the meter.
+        """
         self.val = 0
         self.avg = 0
         self.sum = 0
         self.count = 0
 
     def update(self, val, n=1):
+        """
+        Update with value and weight.
+
+        Parameters
+        ----------
+        val : float or int
+            The new value to be accounted in.
+        n : int
+            The weight of the new value.
+        """
         if not isinstance(val, float) and not isinstance(val, int):
             _logger.warning("Values passed to AverageMeter must be number, not %s.", type(val))
         self.val = val
@@ -112,6 +139,11 @@ class StructuredMutableTreeNode:
     This tree can be seen as a "flattened" version of the module tree. Since nested mutable entity is not supported yet,
     the following must be true: each subtree corresponds to a ``MutableScope`` and each leaf corresponds to a
     ``Mutable`` (other than ``MutableScope``).
+
+    Parameters
+    ----------
+    mutable : nni.nas.pytorch.mutables.Mutable
+        The mutable that current node is linked with.
     """
 
     def __init__(self, mutable):
@@ -119,10 +151,16 @@ class StructuredMutableTreeNode:
         self.children = []
 
     def add_child(self, mutable):
+        """
+        Add a tree node to the children list of current node.
+        """
         self.children.append(StructuredMutableTreeNode(mutable))
         return self.children[-1]
 
     def type(self):
+        """
+        Return the ``type`` of mutable content.
+        """
         return type(self.mutable)
 
     def __iter__(self):