Merge pull request #59 from sony/feature/20221006-add-ofa-related-api…

…-doc

Add the API information of OFASearcher and OFA-based search spaces to docs

docs/source/nnablanas_api.rst

@@ -168,6 +168,13 @@ FairNasSearcher
    :undoc-members:
    :show-inheritance:
 
+OFASearcher
+...............
+.. automodule:: nnabla_nas.runner.searcher.ofa
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Trainer
 .......
 
@@ -251,6 +258,30 @@ FairNas
 .......
 
 .. automodule:: nnabla_nas.contrib.classification.fairnas
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+OFAMobileNetV3
+.......................
+
+.. automodule:: nnabla_nas.contrib.classification.ofa.networks.ofa_mbv3
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+OFAXception
+.......................
+
+.. automodule:: nnabla_nas.contrib.classification.ofa.networks.ofa_xception
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+OFAResnet50
+.......................
+
+.. automodule:: nnabla_nas.contrib.classification.ofa.networks.ofa_resnet50
    :members:
    :undoc-members:
    :show-inheritance:

nnabla_nas/contrib/classification/ofa/networks/ofa_mbv3.py

@@ -85,29 +85,25 @@ def genotype2subnetlist(op_candidates, genotype):
 
 
 class OFAMbv3Net(ClassificationModel):
-    r""" MobileNet V3 Search Net
-    This implementation is based on the PyTorch implementation.
+    r"""MobileNet V3 Search Net.
 
     Args:
-        num_classes (int): Number of classes
-        bn_param (tuple, optional): BatchNormalization decay rate and eps.
+        num_classes (int, optional): Number of classes. Defaults to 1000.
+        bn_param (tuple, optional): BatchNormalization decay rate and eps. Defaults to (0.9, 1e-5).
         drop_rate (float, optional): Drop rate used in Dropout. Defaults to 0.1.
-        base_stage_width (list of int, optional): A list of base stage
-            channel size. Defaults to None.
-        width_mult (float, optional): Multiplier value to base stage channel size.
-            Defaults to 1.0.
-        op_candidates (str or list of str, optional): Operator choices.
-            Defaults to MB6 3x3.
-        depth_candidates (int or list of int, optional): Depth choices.
-            Defaults to 4.
-        weight (str, optional): The path to weight file. Defaults to
-            None.
+        base_stage_width (list of int, optional): A list of base stage channel size. Defaults to None.
+        width_mult (float, optional): Multiplier value to base stage channel size. Defaults to 1.0.
+        op_candidates (str or list of str, optional): Operator choices. Defaults to "MB6 3x3".
+        depth_candidates (int or list of int, optional): Depth choices. Defaults to 4.
+        compound (bool, optional): Use CompOFA or not. Defaults to False.
+        fixed_kernel (bool, optional): Fix kernel or not. Defaults to False.
+        weight_init (str, optional): Weight initializer. Defaults to 'he_fout'.
+        weights (str, optional): The relative path to weight file. Defaults to None.
 
     References:
-    [1] Cai, Han, et al. "Once-for-all: Train one network and specialize it for
-        efficient deployment." arXiv preprint arXiv:1908.09791 (2019).
+        [1] Cai, Han, et al. "Once-for-all: Train one network and specialize it for
+            efficient deployment." arXiv preprint arXiv:1908.09791 (2019).
     """
-
     CHANNEL_DIVISIBLE = 8
 
     def __init__(self,
@@ -487,21 +483,18 @@ def re_organize_middle_weights(self, expand_ratio_stage=0):
 
 
 class TrainNet(OFAMbv3Net):
-    r""" MobileNet V3 Train Net.
+    r"""MobileNet V3 Train Net.
+
     Args:
-        num_classes (int): Number of classes
-        bn_param (tuple, optional): BatchNormalization decay rate and eps.
+        num_classes (int, optional): Number of classes. Defaults to 1000.
+        bn_param (tuple, optional): BatchNormalization decay rate and eps. Defaults to (0.9, 1e-5).
         drop_rate (float, optional): Drop rate used in Dropout. Defaults to 0.1.
-        base_stage_width (list of int, optional): A list of base stage
-            channel size. Defaults to None.
-        width_mult (float, optional): Multiplier value to base stage channel size.
-            Defaults to 1.0.
-        op_candidates (str or list of str, optional): Operator choices.
-            Defaults to MB6 3x3.
-        depth_candidates (int or list of int, optional): Depth choices.
-            Defaults to 4.
-        weight (str, optional): The path to weight file. Defaults to None.
-        genotype (list of int, optional): A list to operators, Defaults to None.
+        base_stage_width (list of int, optional): A list of base stage channel size. Defaults to None.
+        width_mult (float, optional): Multiplier value to base stage channel size. Defaults to 1.0.
+        op_candidates (str or list of str, optional): Operator choices. Defaults to None.
+        depth_candidates (int or list of int, optional): Depth choices. Defaults to None.
+        genotype (list of int, optional): A list to operators. Defaults to None.
+        weights (str, optional): Relative path to the weights file. Defaults to None.
     """
 
     def __init__(self, num_classes=1000, bn_param=(0.9, 1e-5), drop_rate=0.1,

nnabla_nas/contrib/classification/ofa/networks/ofa_resnet50.py

@@ -32,7 +32,7 @@
 
 
 class OFAResNet50(ClassificationModel):
-    r"""OFAResNet50 Base Class
+    r"""OFAResNet50 Base Class.
 
     This is the Base Class used for both TrainNet and SearchNet.
     This implementation is based on the PyTorch implementation given in References.
@@ -280,6 +280,7 @@ def set_bn_param(self, decay_rate, eps, **kwargs):
 
     def loss(self, outputs, targets, loss_weights=None):
         r"""Return loss computed from a list of outputs and list of targets.
+
         Args:
             outputs (list of nn.Variable):
                 A list of output variables computed from the model.
@@ -290,6 +291,7 @@ def loss(self, outputs, targets, loss_weights=None):
                 contributions of different model outputs.
                 It is expected to have a 1:1 mapping to model outputs.
                 Defaults to None.
+
         Returns:
             nn.Variable: A scalar NNabla Variable represents the loss.
         """
@@ -303,9 +305,11 @@ def kd_loss(self, outputs, logits, targets, loss_weights=None):
 
     def get_net_parameters(self, grad_only=False):
         r"""Returns an `OrderedDict` containing architecture parameters.
+
         Args:
             grad_only (bool, optional): If sets to `True`, then only parameters
                 with `need_grad=True` are returned. Defaults to False.
+
         Returns:
             OrderedDict: A dictionary containing parameters.
         """

nnabla_nas/contrib/classification/ofa/networks/ofa_xception.py

@@ -29,6 +29,21 @@
 from ....common.ofa.utils.common_tools import cross_entropy_loss_with_soft_target
 
 
+def _build_candidates_table():
+    kernel_search_space = [3, 5, 7]
+    depth_search_space = [1, 2, 3]
+    expand_ratio_search_space = [0.6, 0.8, 1]
+    candidates_table = {}
+    for cur_kernel in kernel_search_space:
+        for cur_depth in depth_search_space:
+            for cur_expand_ratio in expand_ratio_search_space:
+                key = f'XP{cur_expand_ratio} {cur_kernel}x{cur_kernel} {cur_depth}'
+                value = {'ks': cur_kernel, 'depth': cur_depth,
+                         'expand_ratio': cur_expand_ratio}
+                candidates_table[key] = value
+    return candidates_table
+
+
 class ProcessGenotype:
 
     r""" ProcessGenotype
@@ -37,29 +52,15 @@ class ProcessGenotype:
     to process the genotypes and op_candidates to get the subnet
     architecture or the search space.
 
-    Operator candidates: "XP{E} {K}x{K} {D}"
-                         E=expand_ratio
-                         K=kernel_size
-                         D=depth_of_block
+    Operator candidates: "XP{E} {K}x{K} {D}", E=expand_ratio, K=kernel_size, D=depth_of_block
 
     Note: If depth of a block==1, expand_ratio will be ignored since we
-          just need in_channels and out_channels for a block with a single
-          layer. So blocks: ["XP0.6 KxK 1", "XP0.8 KxK 1", "XP1 KxK 1"]
-          are equivalent in this architecture design.
+    just need in_channels and out_channels for a block with a single
+    layer. So blocks: ["XP0.6 KxK 1", "XP0.8 KxK 1", "XP1 KxK 1"]
+    are equivalent in this architecture design.
     """
 
-    CANDIDATES = {}
-    KERNEL_SEARCH_SPACE = [3, 5, 7]
-    DEPTH_SEARCH_SPACE = [1, 2, 3]
-    EXPAND_RATIO_SEARCH_SPACE = [0.6, 0.8, 1]
-
-    for cur_kernel in KERNEL_SEARCH_SPACE:
-        for cur_depth in DEPTH_SEARCH_SPACE:
-            for cur_expand_ratio in EXPAND_RATIO_SEARCH_SPACE:
-                key = f'XP{cur_expand_ratio} {cur_kernel}x{cur_kernel} {cur_depth}'
-                value = {'ks': cur_kernel, 'depth': cur_depth,
-                         'expand_ratio': cur_expand_ratio}
-                CANDIDATES[key] = value
+    CANDIDATES = _build_candidates_table()
 
     @classmethod
     def get_search_space(cls, candidates):
@@ -352,6 +353,7 @@ def get_net_parameters(self, grad_only=False):
         Args:
             grad_only (bool, optional): If sets to `True`, then only parameters
                 with `need_grad=True` are returned. Defaults to False.
+
         Returns:
             OrderedDict: A dictionary containing parameters.
         """
@@ -360,9 +362,11 @@ def get_net_parameters(self, grad_only=False):
 
     def get_arch_parameters(self, grad_only=False):
         r"""Returns an `OrderedDict` containing architecture parameters.
+
         Args:
             grad_only (bool, optional): If sets to `True`, then only parameters
                 with `need_grad=True` are returned. Defaults to False.
+
         Returns:
             OrderedDict: A dictionary containing parameters.
         """