diff --git a/src/nni_manager/rest_server/restValidationSchemas.ts b/src/nni_manager/rest_server/restValidationSchemas.ts index f794df4d70..6a61c209c3 100644 --- a/src/nni_manager/rest_server/restValidationSchemas.ts +++ b/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), diff --git a/src/sdk/pynni/nni/__init__.py b/src/sdk/pynni/nni/__init__.py index 0358b45bbf..8465846d4c 100644 --- a/src/sdk/pynni/nni/__init__.py +++ b/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): diff --git a/src/sdk/pynni/nni/nas_utils.py b/src/sdk/pynni/nni/nas_utils.py new file mode 100644 index 0000000000..baca7cf30c --- /dev/null +++ b/src/sdk/pynni/nni/nas_utils.py @@ -0,0 +1,160 @@ +# 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) diff --git a/src/sdk/pynni/nni/smartparam.py b/src/sdk/pynni/nni/smartparam.py index 175df204cd..07519b69ea 100644 --- a/src/sdk/pynni/nni/smartparam.py +++ b/src/sdk/pynni/nni/smartparam.py @@ -23,6 +23,7 @@ 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 @@ def mutable_layer( 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 @@ 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]) - - 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: diff --git a/src/sdk/pynni/tests/test_smartparam.py b/src/sdk/pynni/tests/test_smartparam.py index 33fb783afc..001bd35032 100644 --- a/src/sdk/pynni/tests/test_smartparam.py +++ b/src/sdk/pynni/tests/test_smartparam.py @@ -38,7 +38,13 @@ def setUp(self): '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 @@ def test_lambda_func(self): 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() diff --git a/tools/nni_annotation/__init__.py b/tools/nni_annotation/__init__.py index 20ac69f659..ebfe73f5ba 100644 --- a/tools/nni_annotation/__init__.py +++ b/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 diff --git a/tools/nni_annotation/code_generator.py b/tools/nni_annotation/code_generator.py index 4cc2f1f261..871308e5d4 100644 --- a/tools/nni_annotation/code_generator.py +++ b/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 @@ def visit_Call(self, node): # pylint: disable=invalid-name 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 @@ 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: + # 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 @@ def _visit_string(self, node): 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 @@ def _visit_children(self, node): 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) diff --git a/tools/nni_cmd/config_schema.py b/tools/nni_cmd/config_schema.py index ef5173d8b9..84c7ad3bb5 100644 --- a/tools/nni_cmd/config_schema.py +++ b/tools/nni_cmd/config_schema.py @@ -180,7 +180,8 @@ def setPathCheck(key): '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') } } diff --git a/tools/nni_cmd/launcher.py b/tools/nni_cmd/launcher.py index 8e0eb57fc6..8eaca55fc2 100644 --- a/tools/nni_cmd/launcher.py +++ b/tools/nni_cmd/launcher.py @@ -380,7 +380,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)