Implement enas-mode and oneshot-mode for NAS interface (#1201)

* update

* dev-tf-master

* fix bugs

* fix bugs

* remove unnecessary lines

* dev oneshot and darts

* dev nas

* dev enas and oneshot

* dev enas and oneshot

* dev enas and oneshot

* dev oneshot and enas

* dev oneshot

* add ut

* add docstring

* add docstring

* fix

* resolve comments by changing docstring

* resolve comments

src/nni_manager/rest_server/restValidationSchemas.ts

@@ -51,6 +51,7 @@ export namespace ValidationSchemas {
                 command: joi.string().min(1),
                 virtualCluster: joi.string(),
                 shmMB: joi.number(),
+                nasMode: joi.string().valid('classic_mode', 'enas_mode', 'oneshot_mode'),
                 worker: joi.object({
                     replicas: joi.number().min(1).required(),
                     image: joi.string().min(1),

src/sdk/pynni/nni/__init__.py

@@ -23,6 +23,7 @@
 
 from .trial import *
 from .smartparam import *
+from .nas_utils import reload_tensorflow_variables
 
 class NoMoreTrialError(Exception):
     def __init__(self,ErrorInfo):

src/sdk/pynni/nni/nas_utils.py

+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
+# associated documentation files (the "Software"), to deal in the Software without restriction,
+# including without limitation the rights to use, copy, modify, merge, publish, distribute,
+# sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or
+# substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
+# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT
+# OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+# ==================================================================================================
+
+from . import trial
+
+
+def classic_mode(
+        mutable_id,
+        mutable_layer_id,
+        funcs,
+        funcs_args,
+        fixed_inputs,
+        optional_inputs,
+        optional_input_size):
+    '''Execute the chosen function and inputs directly.
+    In this mode, the trial code is only running the chosen subgraph (i.e., the chosen ops and inputs),
+    without touching the full model graph.'''
+    if trial._params is None:
+        trial.get_next_parameter()
+    mutable_block = trial.get_current_parameter(mutable_id)
+    chosen_layer = mutable_block[mutable_layer_id]["chosen_layer"]
+    chosen_inputs = mutable_block[mutable_layer_id]["chosen_inputs"]
+    real_chosen_inputs = [optional_inputs[input_name]
+                          for input_name in chosen_inputs]
+    layer_out = funcs[chosen_layer](
+        [fixed_inputs, real_chosen_inputs], **funcs_args[chosen_layer])
+
+    return layer_out
+
+
+def enas_mode(
+        mutable_id,
+        mutable_layer_id,
+        funcs,
+        funcs_args,
+        fixed_inputs,
+        optional_inputs,
+        optional_input_size,
+        tf):
+    '''For enas mode, we build the full model graph in trial but only run a subgraph。
+    This is implemented by masking inputs and branching ops.
+    Specifically, based on the received subgraph (through nni.get_next_parameter),
+    it can be known which inputs should be masked and which op should be executed.'''
+    name_prefix = "{}_{}".format(mutable_id, mutable_layer_id)
+    # store namespace
+    if 'name_space' not in globals():
+        global name_space
+        name_space = dict()
+    name_space[mutable_id] = True
+    name_space[name_prefix] = dict()
+    name_space[name_prefix]['funcs'] = list(funcs)
+    name_space[name_prefix]['optional_inputs'] = list(optional_inputs)
+    # create tensorflow variables as 1/0 signals used to form subgraph
+    if 'tf_variables' not in globals():
+        global tf_variables
+        tf_variables = dict()
+    name_for_optional_inputs = name_prefix + '_optional_inputs'
+    name_for_funcs = name_prefix + '_funcs'
+    tf_variables[name_prefix] = dict()
+    tf_variables[name_prefix]['optional_inputs'] = tf.get_variable(name_for_optional_inputs,
+                                                                   [len(
+                                                                       optional_inputs)],
+                                                                   dtype=tf.bool,
+                                                                   trainable=False)
+    tf_variables[name_prefix]['funcs'] = tf.get_variable(
+        name_for_funcs, [], dtype=tf.int64, trainable=False)
+
+    # get real values using their variable names
+    real_optional_inputs_value = [optional_inputs[name]
+                                  for name in name_space[name_prefix]['optional_inputs']]
+    real_func_value = [funcs[name]
+                       for name in name_space[name_prefix]['funcs']]
+    real_funcs_args = [funcs_args[name]
+                       for name in name_space[name_prefix]['funcs']]
+    # build tensorflow graph of geting chosen inputs by masking
+    real_chosen_inputs = tf.boolean_mask(
+        real_optional_inputs_value, tf_variables[name_prefix]['optional_inputs'])
+    # build tensorflow graph of different branches by using tf.case
+    branches = dict()
+    for func_id in range(len(funcs)):
+        func_output = real_func_value[func_id](
+            [fixed_inputs, real_chosen_inputs], **real_funcs_args[func_id])
+        branches[tf.equal(tf_variables[name_prefix]['funcs'],
+                          func_id)] = lambda: func_output
+    layer_out = tf.case(branches, exclusive=True,
+                        default=lambda: func_output)
+
+    return layer_out
+
+
+def oneshot_mode(
+        mutable_id,
+        mutable_layer_id,
+        funcs,
+        funcs_args,
+        fixed_inputs,
+        optional_inputs,
+        optional_input_size,
+        tf):
+    '''Similar to enas mode, oneshot mode also builds the full model graph.
+    The difference is that oneshot mode does not receive subgraph.
+    Instead, it uses dropout to randomly dropout inputs and ops.'''
+    # NNI requires to get_next_parameter before report a result. But the parameter will not be used in this mode
+    if trial._params is None:
+        trial.get_next_parameter()
+    optional_inputs = list(optional_inputs.values())
+    inputs_num = len(optional_inputs)
+    # Calculate dropout rate according to the formular r^(1/k), where r is a hyper-parameter and k is the number of inputs
+    if inputs_num > 0:
+        rate = 0.01 ** (1 / inputs_num)
+        noise_shape = [inputs_num] + [1] * len(optional_inputs[0].get_shape())
+        optional_inputs = tf.nn.dropout(
+            optional_inputs, rate=rate, noise_shape=noise_shape)
+        optional_inputs = [optional_inputs[idx] for idx in range(inputs_num)]
+    layer_outs = [func([fixed_inputs, optional_inputs], **funcs_args[func_name])
+                  for func_name, func in funcs.items()]
+    layer_out = tf.add_n(layer_outs)
+
+    return layer_out
+
+
+def reload_tensorflow_variables(session, tf=None):
+    '''In Enas mode, this function reload every signal varaible created in `enas_mode` function so
+    the whole tensorflow graph will be changed into certain subgraph recerived from Tuner.
+    ---------------
+    session: the tensorflow session created by users
+    tf: tensorflow module
+    '''
+    subgraph_from_tuner = trial.get_next_parameter()
+    for mutable_id, mutable_block in subgraph_from_tuner.items():
+        if mutable_id not in name_space:
+            continue
+        for mutable_layer_id, mutable_layer in mutable_block.items():
+            name_prefix = "{}_{}".format(mutable_id, mutable_layer_id)
+            # extract layer information from the subgraph sampled by tuner
+            chosen_layer = name_space[name_prefix]['funcs'].index(
+                mutable_layer["chosen_layer"])
+            chosen_inputs = [1 if inp in mutable_layer["chosen_inputs"]
+                             else 0 for inp in name_space[name_prefix]['optional_inputs']]
+            # load these information into pre-defined tensorflow variables
+            tf_variables[name_prefix]['funcs'].load(chosen_layer, session)
+            tf_variables[name_prefix]['optional_inputs'].load(
+                chosen_inputs, session)

src/sdk/pynni/nni/smartparam.py

@@ -23,6 +23,7 @@ import random
 
 from .env_vars import trial_env_vars
 from . import trial
+from .nas_utils import classic_mode, enas_mode, oneshot_mode
 
 
 __all__ = [
@@ -124,7 +125,9 @@ else:
             funcs_args,
             fixed_inputs,
             optional_inputs,
-            optional_input_size):
+            optional_input_size,
+            mode='classic_mode',
+            tf=None):
         '''execute the chosen function and inputs.
         Below is an example of chosen function and inputs:
         {
@@ -144,14 +147,38 @@ else:
         fixed_inputs:
         optional_inputs: dict of optional inputs
         optional_input_size: number of candidate inputs to be chosen
+        tf: tensorflow module
         '''
-        mutable_block = _get_param(mutable_id)
-        chosen_layer = mutable_block[mutable_layer_id]["chosen_layer"]
-        chosen_inputs = mutable_block[mutable_layer_id]["chosen_inputs"]
-        real_chosen_inputs = [optional_inputs[input_name] for input_name in chosen_inputs]
-        layer_out = funcs[chosen_layer]([fixed_inputs, real_chosen_inputs], **funcs_args[chosen_layer])
-        
-        return layer_out
+        if mode == 'classic_mode':
+            return classic_mode(mutable_id,
+                            mutable_layer_id,
+                            funcs,
+                            funcs_args,
+                            fixed_inputs,
+                            optional_inputs,
+                            optional_input_size)
+        elif mode == 'enas_mode':
+            assert tf is not None, 'Internal Error: Tensorflow should not be None in enas_mode'
+            return enas_mode(mutable_id,
+                            mutable_layer_id,
+                            funcs,
+                            funcs_args,
+                            fixed_inputs,
+                            optional_inputs,
+                            optional_input_size,
+                            tf)
+        elif mode == 'oneshot_mode':
+            assert tf is not None, 'Internal Error: Tensorflow should not be None in oneshot_mode'
+            return oneshot_mode(mutable_id,
+                            mutable_layer_id,
+                            funcs,
+                            funcs_args,
+                            fixed_inputs,
+                            optional_inputs,
+                            optional_input_size,
+                            tf)
+        else:
+            raise RuntimeError('Unrecognized mode: %s' % mode)
 
     def _get_param(key):
         if trial._params is None:

src/sdk/pynni/tests/test_smartparam.py

@@ -38,7 +38,13 @@ class SmartParamTestCase(TestCase):
             'test_smartparam/choice3/choice': '[1, 2]',
             'test_smartparam/choice4/choice': '{"a", 2}',
             'test_smartparam/func/function_choice': 'bar',
-            'test_smartparam/lambda_func/function_choice': "lambda: 2*3"
+            'test_smartparam/lambda_func/function_choice': "lambda: 2*3",
+            'mutable_block_66':{
+                'mutable_layer_0':{
+                    'chosen_layer': 'conv2D(size=5)',
+                    'chosen_inputs': ['y']
+                }
+            }
         }
         nni.trial._params = { 'parameter_id': 'test_trial', 'parameters': params }
 
@@ -61,6 +67,13 @@ class SmartParamTestCase(TestCase):
         val = nni.function_choice({"lambda: 2*3": lambda: 2*3, "lambda: 3*4": lambda: 3*4}, name = 'lambda_func', key='test_smartparam/lambda_func/function_choice')
         self.assertEqual(val, 6)
 
+    def test_mutable_layer(self):
+        layer_out = nni.mutable_layer('mutable_block_66',
+                'mutable_layer_0', {'conv2D(size=3)': conv2D, 'conv2D(size=5)': conv2D}, {'conv2D(size=3)':
+                {'size':3}, 'conv2D(size=5)': {'size':5}}, [100], {'x':1,'y':2}, 1, 'classic_mode')
+        self.assertEqual(layer_out, [100, 2, 5])
+        
+
 
 def foo():
     return 'foo'
@@ -68,6 +81,8 @@ def foo():
 def bar():
     return 'bar'
 
+def conv2D(inputs, size=3):
+    return inputs[0] + inputs[1] + [size]
 
 if __name__ == '__main__':
     main()

tools/nni_annotation/__init__.py

@@ -76,11 +76,12 @@ def _generate_file_search_space(path, module):
     return search_space
 
 
-def expand_annotations(src_dir, dst_dir, exp_id='', trial_id=''):
+def expand_annotations(src_dir, dst_dir, exp_id='', trial_id='', nas_mode=None):
     """Expand annotations in user code.
     Return dst_dir if annotation detected; return src_dir if not.
     src_dir: directory path of user code (str)
     dst_dir: directory to place generated files (str)
+    nas_mode: the mode of NAS given that NAS interface is used
     """
     if src_dir[-1] == slash:
         src_dir = src_dir[:-1]
@@ -108,7 +109,7 @@ def expand_annotations(src_dir, dst_dir, exp_id='', trial_id=''):
             dst_path = os.path.join(dst_subdir, file_name)
             if file_name.endswith('.py'):
                 if trial_id == '':
-                    annotated |= _expand_file_annotations(src_path, dst_path)
+                    annotated |= _expand_file_annotations(src_path, dst_path, nas_mode)
                 else:
                     module = package + file_name[:-3]
                     annotated |= _generate_specific_file(src_path, dst_path, exp_id, trial_id, module)
@@ -120,10 +121,10 @@ def expand_annotations(src_dir, dst_dir, exp_id='', trial_id=''):
 
     return dst_dir if annotated else src_dir
 
-def _expand_file_annotations(src_path, dst_path):
+def _expand_file_annotations(src_path, dst_path, nas_mode):
     with open(src_path) as src, open(dst_path, 'w') as dst:
         try:
-            annotated_code = code_generator.parse(src.read())
+            annotated_code = code_generator.parse(src.read(), nas_mode)
             if annotated_code is None:
                 shutil.copyfile(src_path, dst_path)
                 return False

tools/nni_annotation/code_generator.py

@@ -21,14 +21,14 @@
 
 import ast
 import astor
-from nni_cmd.common_utils import print_warning
 
 # pylint: disable=unidiomatic-typecheck
 
-def parse_annotation_mutable_layers(code, lineno):
+def parse_annotation_mutable_layers(code, lineno, nas_mode):
     """Parse the string of mutable layers in annotation.
     Return a list of AST Expr nodes
     code: annotation string (excluding '@')
+    nas_mode: the mode of NAS
     """
     module = ast.parse(code)
     assert type(module) is ast.Module, 'internal error #1'
@@ -110,6 +110,9 @@ def parse_annotation_mutable_layers(code, lineno):
         else:
             target_call_args.append(ast.Dict(keys=[], values=[]))
             target_call_args.append(ast.Num(n=0))
+        target_call_args.append(ast.Str(s=nas_mode))
+        if nas_mode in ['enas_mode', 'oneshot_mode']:
+            target_call_args.append(ast.Name(id='tensorflow'))
         target_call = ast.Call(func=target_call_attr, args=target_call_args, keywords=[])
         node = ast.Assign(targets=[layer_output], value=target_call)
         nodes.append(node)
@@ -277,10 +280,11 @@ class FuncReplacer(ast.NodeTransformer):
 class Transformer(ast.NodeTransformer):
     """Transform original code to annotated code"""
 
-    def __init__(self):
+    def __init__(self, nas_mode=None):
         self.stack = []
         self.last_line = 0
         self.annotated = False
+        self.nas_mode = nas_mode
 
     def visit(self, node):
         if isinstance(node, (ast.expr, ast.stmt)):
@@ -316,8 +320,11 @@ class Transformer(ast.NodeTransformer):
             return node  # not an annotation, ignore it
 
         if string.startswith('@nni.get_next_parameter'):
-            deprecated_message = "'@nni.get_next_parameter' is deprecated in annotation due to inconvenience. Please remove this line in the trial code."
-            print_warning(deprecated_message)
+            call_node = parse_annotation(string[1:]).value
+            if call_node.args:
+                # it is used in enas mode as it needs to retrieve the next subgraph for training
+                call_attr = ast.Attribute(value=ast.Name(id='nni', ctx=ast.Load()), attr='reload_tensorflow_variables', ctx=ast.Load())
+                return ast.Expr(value=ast.Call(func=call_attr, args=call_node.args, keywords=[]))
 
         if string.startswith('@nni.report_intermediate_result')  \
                 or string.startswith('@nni.report_final_result') \
@@ -325,7 +332,8 @@ class Transformer(ast.NodeTransformer):
             return parse_annotation(string[1:])  # expand annotation string to code
 
         if string.startswith('@nni.mutable_layers'):
-            return parse_annotation_mutable_layers(string[1:], node.lineno)
+            nodes = parse_annotation_mutable_layers(string[1:], node.lineno, self.nas_mode)
+            return nodes
 
         if string.startswith('@nni.variable') \
                 or string.startswith('@nni.function_choice'):
@@ -343,17 +351,18 @@ class Transformer(ast.NodeTransformer):
         return node
 
 
-def parse(code):
+def parse(code, nas_mode=None):
     """Annotate user code.
     Return annotated code (str) if annotation detected; return None if not.
-    code: original user code (str)
+    code: original user code (str),
+    nas_mode: the mode of NAS given that NAS interface is used
     """
     try:
         ast_tree = ast.parse(code)
     except Exception:
         raise RuntimeError('Bad Python code')
 
-    transformer = Transformer()
+    transformer = Transformer(nas_mode)
     try:
         transformer.visit(ast_tree)
     except AssertionError as exc:
@@ -369,5 +378,9 @@ def parse(code):
         if type(nodes[i]) is ast.ImportFrom and nodes[i].module == '__future__':
             last_future_import = i
     nodes.insert(last_future_import + 1, import_nni)
+    # enas and oneshot modes for tensorflow need tensorflow module, so we import it here
+    if nas_mode in ['enas_mode', 'oneshot_mode']:
+        import_tf = ast.Import(names=[ast.alias(name='tensorflow', asname=None)])
+        nodes.insert(last_future_import + 1, import_tf)
 
     return astor.to_source(ast_tree)

tools/nni_cmd/config_schema.py

@@ -181,7 +181,8 @@ common_trial_schema = {
 'trial':{
     'command': setType('command', str),
     'codeDir': setPathCheck('codeDir'),
-    'gpuNum': setNumberRange('gpuNum', int, 0, 99999)
+    'gpuNum': setNumberRange('gpuNum', int, 0, 99999),
+    Optional('nasMode'): setChoice('classic_mode', 'enas_mode', 'oneshot_mode')
     }
 }
 

tools/nni_cmd/launcher.py

@@ -377,7 +377,7 @@ def launch_experiment(args, experiment_config, mode, config_file_name, experimen
         if not os.path.isdir(path):
             os.makedirs(path)
         path = tempfile.mkdtemp(dir=path)
-        code_dir = expand_annotations(experiment_config['trial']['codeDir'], path)
+        code_dir = expand_annotations(experiment_config['trial']['codeDir'], path, nas_mode=experiment_config['trial']['nasMode'])
         experiment_config['trial']['codeDir'] = code_dir
         search_space = generate_search_space(code_dir)
         experiment_config['searchSpace'] = json.dumps(search_space)