Implement full graph version (for tensorflow) of NAS-Interface

src/sdk/pynni/nni/smartparam.py

@@ -37,7 +37,8 @@
     'lognormal',
     'qlognormal',
     'function_choice',
-    'mutable_layer'
+    'mutable_layer',
+    'reload_tensorflow_variables'
 ]
 
 
@@ -82,6 +83,10 @@ def function_choice(*funcs, name=None):
     def mutable_layer():
         raise RuntimeError('Cannot call nni.mutable_layer in this mode')
 
+    def reload_tensorflow_variables(session, tf):
+        raise RuntimeError(
+            'Cannot call nni.reload_tensorflow_variables in this mode')
+
 else:
 
     def choice(options, name=None, key=None):
@@ -124,7 +129,8 @@ def mutable_layer(
             funcs_args,
             fixed_inputs,
             optional_inputs,
-            optional_input_size):
+            optional_input_size,
+            tf=None):
         '''execute the chosen function and inputs.
         Below is an example of chosen function and inputs:
         {
@@ -144,15 +150,77 @@ def mutable_layer(
         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])
-
+        if tf is None:
+            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])
+        else:
+            name_prefix = "{}_{}".format(mutable_id, mutable_layer_id)
+            # store namespace
+            global name_space
+            if 'name_space' not in globals():
+                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
+            global tf_variables
+            if 'tf_variables' not in globals():
+                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 reload_tensorflow_variables(session):
+        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)
+
     def _get_param(key):
         if trial._params is None:
             trial.get_next_parameter()

tools/nni_annotation/code_generator.py

@@ -21,13 +21,12 @@
 
 import ast
 import astor
-from nni_cmd.common_utils import print_warning
 
 # pylint: disable=unidiomatic-typecheck
 
 def parse_annotation_mutable_layers(code, lineno):
     """Parse the string of mutable layers in annotation.
-    Return a list of AST Expr nodes
+    Return a list of AST Expr nodes and NAS mode
     code: annotation string (excluding '@')
     """
     module = ast.parse(code)
@@ -38,7 +37,17 @@ def parse_annotation_mutable_layers(code, lineno):
     nodes = []
     mutable_id = 'mutable_block_' + str(lineno)
     mutable_layer_cnt = 0
+    if call.keywords:
+        mode = call.keywords[0].value
+        assert isinstance(mode, (ast.Str, ast.Name)), 'Mode must be a string or Name'
+        mode = mode.s if isinstance(mode, ast.Str) else mode.id
+        assert mode in ['general', 'oneshot-tf', 'oneshot-pytorch']
+    else:
+        mode = 'general'
     for arg in call.args:
+        if type(arg) is ast.Str:
+            assert arg.s in ['general', 'tensorflow'], 'Unsupported mode %s' % arg.s
+            mode = arg.s
         fields = {'layer_choice': False,
                   'fixed_inputs': False,
                   'optional_inputs': False,
@@ -110,10 +119,12 @@ def parse_annotation_mutable_layers(code, lineno):
         else:
             target_call_args.append(ast.Dict(keys=[], values=[]))
             target_call_args.append(ast.Num(n=0))
+        if mode == 'oneshot-tf':
+            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)
-    return nodes
+    return nodes, mode
 
 def parse_annotation(code):
     """Parse an annotation string.
@@ -281,6 +292,7 @@ def __init__(self):
         self.stack = []
         self.last_line = 0
         self.annotated = False
+        self.mode = 'general'
 
     def visit(self, node):
         if isinstance(node, (ast.expr, ast.stmt)):
@@ -316,16 +328,19 @@ def _visit_string(self, node):
             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:
+                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') \
                 or string.startswith('@nni.get_next_parameter'):
             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, self.mode = parse_annotation_mutable_layers(string[1:], node.lineno)
+            return nodes
 
         if string.startswith('@nni.variable') \
                 or string.startswith('@nni.function_choice'):
@@ -369,5 +384,8 @@ 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)
+    if transformer.mode == 'oneshot-tf':
+        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)