Update oneshot_mode and add darts mode (first version) (#1369)

* dev oneshot and darts mode, add mnist example

azure-pipelines.yml

@@ -15,7 +15,7 @@ jobs:
     displayName: 'Install nni toolkit via source code'
   - script: |
       python3 -m pip install flake8 --user
-      IGNORE=./tools/nni_annotation/testcase/*:F821,./examples/trials/mnist-nas/mnist.py:F821
+      IGNORE=./tools/nni_annotation/testcase/*:F821,./examples/trials/mnist-nas/*/mnist*.py:F821
       python3 -m flake8 . --count --per-file-ignores=$IGNORE --select=E9,F63,F72,F82 --show-source --statistics
     displayName: 'Run flake8 tests to find Python syntax errors and undefined names'
   - script: |

examples/trials/mnist-nas/config.yml → examples/trials/mnist-nas/classic_mode/config_nas.yml

@@ -8,13 +8,11 @@ trainingServicePlatform: local
 #choice: true, false
 useAnnotation: true
 tuner:
-  #choice: TPE, Random, Anneal, Evolution, BatchTuner, MetisTuner
-  #SMAC (SMAC should be installed through nnictl)
-  #codeDir: ~/nni/nni/examples/tuners/random_nas_tuner
-  codeDir: ../../tuners/random_nas_tuner
+  codeDir: ../../../tuners/random_nas_tuner
   classFileName: random_nas_tuner.py
   className: RandomNASTuner
 trial:
   command: python3 mnist.py
   codeDir: .
   gpuNum: 0
+  nasMode: classic_mode

examples/trials/mnist-nas/darts_mode/config_darts.yml

+authorName: default
+experimentName: example_mnist
+trialConcurrency: 1
+maxExecDuration: 1h
+maxTrialNum: 10
+#choice: local, remote, pai
+trainingServicePlatform: local
+#choice: true, false
+useAnnotation: true
+tuner:
+  codeDir: ../../../tuners/random_nas_tuner
+  classFileName: random_nas_tuner.py
+  className: RandomNASTuner
+trial:
+  command: python3 mnist-darts.py
+  codeDir: .
+  gpuNum: 0
+  nasMode: oneshot_mode

examples/trials/mnist-nas/darts_mode/mnist-darts.py

+"""A deep MNIST classifier using convolutional layers."""
+
+import argparse
+import logging
+import math
+import tempfile
+import time
+
+import tensorflow as tf
+from tensorflow.examples.tutorials.mnist import input_data
+
+import operators as op
+
+FLAGS = None
+
+logger = logging.getLogger('mnist_AutoML')
+
+
+class MnistNetwork(object):
+    '''
+    MnistNetwork is for initializing and building basic network for mnist.
+    '''
+    def __init__(self,
+                 channel_1_num,
+                 channel_2_num,
+                 conv_size,
+                 hidden_size,
+                 pool_size,
+                 learning_rate,
+                 x_dim=784,
+                 y_dim=10):
+        self.channel_1_num = channel_1_num
+        self.channel_2_num = channel_2_num
+        self.conv_size = conv_size
+        self.hidden_size = hidden_size
+        self.pool_size = pool_size
+        self.learning_rate = learning_rate
+        self.x_dim = x_dim
+        self.y_dim = y_dim
+
+        self.images = tf.placeholder(tf.float32, [None, self.x_dim], name='input_x')
+        self.labels = tf.placeholder(tf.float32, [None, self.y_dim], name='input_y')
+        self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')
+
+        self.train_step = None
+        self.accuracy = None
+
+    def build_network(self):
+        '''
+        Building network for mnist, meanwhile specifying its neural architecture search space
+        '''
+
+        # Reshape to use within a convolutional neural net.
+        # Last dimension is for "features" - there is only one here, since images are
+        # grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.
+        with tf.name_scope('reshape'):
+            try:
+                input_dim = int(math.sqrt(self.x_dim))
+            except:
+                print(
+                    'input dim cannot be sqrt and reshape. input dim: ' + str(self.x_dim))
+                logger.debug(
+                    'input dim cannot be sqrt and reshape. input dim: %s', str(self.x_dim))
+                raise
+            x_image = tf.reshape(self.images, [-1, input_dim, input_dim, 1])
+
+        """@nni.mutable_layers(
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=1, out_ch=self.channel_1_num),
+                               op.conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=7, in_ch=1, out_ch=self.channel_1_num),
+                               op.dilated_conv(in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=5, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=7, in_ch=1, out_ch=self.channel_1_num)],
+                fixed_inputs: [x_image],
+                layer_output: conv1_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_1_num)],
+                fixed_inputs: [conv1_out],
+                layer_output: post1_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post1_out],
+                layer_output: pool1_out
+            },
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.dilated_conv(in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=5, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num)],
+                fixed_inputs: [pool1_out],
+                optional_inputs: [post1_out],
+                optional_input_size: [0, 1],
+                layer_output: conv2_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_2_num)],
+                fixed_inputs: [conv2_out],
+                layer_output: post2_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post2_out],
+                optional_inputs: [post1_out, pool1_out],
+                optional_input_size: [0, 1],
+                layer_output: pool2_out
+            }
+        )"""
+
+        # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image
+        # is down to 7x7x64 feature maps -- maps this to 1024 features.
+        last_dim_list = pool2_out.get_shape().as_list()
+        assert(last_dim_list[1] == last_dim_list[2])
+        last_dim = last_dim_list[1]
+        with tf.name_scope('fc1'):
+            w_fc1 = op.weight_variable(
+                [last_dim * last_dim * self.channel_2_num, self.hidden_size])
+            b_fc1 = op.bias_variable([self.hidden_size])
+
+        h_pool2_flat = tf.reshape(
+            pool2_out, [-1, last_dim * last_dim * self.channel_2_num])
+        h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)
+
+        # Dropout - controls the complexity of the model, prevents co-adaptation of features.
+        with tf.name_scope('dropout'):
+            h_fc1_drop = tf.nn.dropout(h_fc1, self.keep_prob)
+
+        # Map the 1024 features to 10 classes, one for each digit
+        with tf.name_scope('fc2'):
+            w_fc2 = op.weight_variable([self.hidden_size, self.y_dim])
+            b_fc2 = op.bias_variable([self.y_dim])
+            y_conv = tf.matmul(h_fc1_drop, w_fc2) + b_fc2
+
+        with tf.name_scope('loss'):
+            self.cross_entropy = tf.reduce_mean(
+                tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=y_conv))
+        with tf.name_scope('adam_optimizer'):
+            self.train_step = tf.train.AdamOptimizer(
+                self.learning_rate).minimize(self.cross_entropy)
+
+        with tf.name_scope('accuracy'):
+            correct_prediction = tf.equal(
+                tf.argmax(y_conv, 1), tf.argmax(self.labels, 1))
+            self.accuracy = tf.reduce_mean(
+                tf.cast(correct_prediction, tf.float32))
+
+
+def download_mnist_retry(data_dir, max_num_retries=20):
+    """Try to download mnist dataset and avoid errors"""
+    for _ in range(max_num_retries):
+        try:
+            return input_data.read_data_sets(data_dir, one_hot=True)
+        except tf.errors.AlreadyExistsError:
+            time.sleep(1)
+    raise Exception("Failed to download MNIST.")
+
+def main(params):
+    '''
+    Main function, build mnist network, run and send result to NNI.
+    '''
+    # Import data
+    mnist = download_mnist_retry(params['data_dir'])
+    print('Mnist download data done.')
+    logger.debug('Mnist download data done.')
+
+    # Create the model
+    # Build the graph for the deep net
+    mnist_network = MnistNetwork(channel_1_num=params['channel_1_num'],
+                                 channel_2_num=params['channel_2_num'],
+                                 conv_size=params['conv_size'],
+                                 hidden_size=params['hidden_size'],
+                                 pool_size=params['pool_size'],
+                                 learning_rate=params['learning_rate'])
+    mnist_network.build_network()
+    logger.debug('Mnist build network done.')
+
+    # Write log
+    graph_location = tempfile.mkdtemp()
+    logger.debug('Saving graph to: %s', graph_location)
+    train_writer = tf.summary.FileWriter(graph_location)
+    train_writer.add_graph(tf.get_default_graph())
+
+    test_acc = 0.0
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        for i in range(params['batch_num']):
+            batch = mnist.train.next_batch(params['batch_size'])
+            feed_dict={mnist_network.images: batch[0],
+                        mnist_network.labels: batch[1],
+                        mnist_network.keep_prob: 1 - params['dropout_rate']}
+            """@nni.training_update(tf, sess, mnist_network.cross_entropy)"""
+            batch = mnist.train.next_batch(params['batch_size'])
+            feed_dict={mnist_network.images: batch[0],
+                        mnist_network.labels: batch[1],
+                        mnist_network.keep_prob: 1 - params['dropout_rate']}
+            mnist_network.train_step.run(feed_dict=feed_dict)
+
+            if i % 100 == 0:
+                test_acc = mnist_network.accuracy.eval(
+                    feed_dict={mnist_network.images: mnist.test.images,
+                               mnist_network.labels: mnist.test.labels,
+                               mnist_network.keep_prob: 1.0})
+
+                """@nni.report_intermediate_result(test_acc)"""
+                logger.debug('test accuracy %g', test_acc)
+                logger.debug('Pipe send intermediate result done.')
+
+        test_acc = mnist_network.accuracy.eval(
+            feed_dict={mnist_network.images: mnist.test.images,
+                       mnist_network.labels: mnist.test.labels,
+                       mnist_network.keep_prob: 1.0})
+
+        """@nni.report_final_result(test_acc)"""
+        logger.debug('Final result is %g', test_acc)
+        logger.debug('Send final result done.')
+
+def get_params():
+    ''' Get parameters from command line '''
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_dir", type=str, default='/tmp/tensorflow/mnist/input_data', help="data directory")
+    parser.add_argument("--dropout_rate", type=float, default=0.5, help="dropout rate")
+    parser.add_argument("--channel_1_num", type=int, default=32)
+    parser.add_argument("--channel_2_num", type=int, default=64)
+    parser.add_argument("--conv_size", type=int, default=5)
+    parser.add_argument("--pool_size", type=int, default=2)
+    parser.add_argument("--hidden_size", type=int, default=1024)
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--batch_num", type=int, default=2000)
+    parser.add_argument("--batch_size", type=int, default=32)
+
+    args, _ = parser.parse_known_args()
+    return args
+
+if __name__ == '__main__':
+    try:
+        params = vars(get_params())
+        main(params)
+    except Exception as exception:
+        logger.exception(exception)
+        raise

examples/trials/mnist-nas/darts_mode/operators.py

+import tensorflow as tf
+import math
+
+
+def weight_variable(shape):
+    """weight_variable generates a weight variable of a given shape."""
+    initial = tf.truncated_normal(shape, stddev=0.1)
+    return tf.Variable(initial)
+
+def bias_variable(shape):
+    """bias_variable generates a bias variable of a given shape."""
+    initial = tf.constant(0.1, shape=shape)
+    return tf.Variable(initial)
+
+def sum_op(inputs):
+    """sum_op"""
+    fixed_input = inputs[0][0]
+    optional_input = tf.concat(inputs[1], axis=3)
+    fixed_shape = fixed_input.get_shape().as_list()
+    optional_shape = optional_input.get_shape().as_list()
+    assert fixed_shape[1] == fixed_shape[2]
+    assert optional_shape[1] == optional_shape[2]
+    pool_size = math.ceil(optional_shape[1] / fixed_shape[1])
+    pool_out = tf.nn.avg_pool(optional_input, ksize=[1, pool_size, pool_size, 1], strides=[1, pool_size, pool_size, 1], padding='SAME')
+    conv_matrix = weight_variable([1, 1, optional_shape[3], fixed_shape[3]])
+    conv_out = tf.nn.conv2d(pool_out, conv_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    return fixed_input + conv_out
+
+
+def conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """conv2d returns a 2d convolution layer with full stride."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [1, 3]:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.conv2d(x_input, w_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def twice_conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """twice_conv2d"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in  [3, 7]:
+        w_matrix1 = weight_variable([1, size, in_ch, int(out_ch/2)])
+        out = tf.nn.conv2d(x_input, w_matrix1, strides=[1, 1, 1, 1], padding='SAME')
+        w_matrix2 = weight_variable([size, 1, int(out_ch/2), out_ch])
+        return tf.nn.conv2d(out, w_matrix2, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def dilated_conv(inputs, size=3, in_ch=-1, out_ch=-1):
+    """dilated_conv"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size == 3:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.atrous_conv2d(x_input, w_matrix, rate=2, padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def separable_conv(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """separable_conv"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        depth_matrix = weight_variable([size, size, in_ch, 1])
+        point_matrix = weight_variable([1, 1, 1*in_ch, out_ch])
+        return tf.nn.separable_conv2d(x_input, depth_matrix, point_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def avg_pool(inputs, size=-1):
+    """avg_pool downsamples a feature map."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.avg_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def max_pool(inputs, size=-1):
+    """max_pool downsamples a feature map."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.max_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def post_process(inputs, ch_size=-1):
+    """post_process"""
+    x_input = inputs[0][0]
+    bias_matrix = bias_variable([ch_size])
+    return tf.nn.relu(x_input + bias_matrix)

examples/trials/mnist-nas/enas_mode/config_enas.yml

+authorName: default
+experimentName: example_mnist
+trialConcurrency: 1
+maxExecDuration: 1h
+maxTrialNum: 10
+#choice: local, remote, pai
+trainingServicePlatform: local
+#choice: true, false
+useAnnotation: true
+multiPhase: true
+tuner:
+  codeDir: ../../../tuners/random_nas_tuner
+  classFileName: random_nas_tuner.py
+  className: RandomNASTuner
+trial:
+  command: python3 mnist-enas.py
+  codeDir: .
+  gpuNum: 0
+  nasMode: enas_mode

examples/trials/mnist-nas/enas_mode/mnist-enas.py

+"""A deep MNIST classifier using convolutional layers."""
+
+import argparse
+import logging
+import math
+import tempfile
+import time
+
+import tensorflow as tf
+from tensorflow.examples.tutorials.mnist import input_data
+
+import operators as op
+
+FLAGS = None
+
+logger = logging.getLogger('mnist_AutoML')
+
+
+class MnistNetwork(object):
+    '''
+    MnistNetwork is for initializing and building basic network for mnist.
+    '''
+    def __init__(self,
+                 channel_1_num,
+                 channel_2_num,
+                 conv_size,
+                 hidden_size,
+                 pool_size,
+                 learning_rate,
+                 x_dim=784,
+                 y_dim=10):
+        self.channel_1_num = channel_1_num
+        self.channel_2_num = channel_2_num
+        self.conv_size = conv_size
+        self.hidden_size = hidden_size
+        self.pool_size = pool_size
+        self.learning_rate = learning_rate
+        self.x_dim = x_dim
+        self.y_dim = y_dim
+
+        self.images = tf.placeholder(tf.float32, [None, self.x_dim], name='input_x')
+        self.labels = tf.placeholder(tf.float32, [None, self.y_dim], name='input_y')
+        self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')
+
+        self.train_step = None
+        self.accuracy = None
+
+    def build_network(self):
+        '''
+        Building network for mnist, meanwhile specifying its neural architecture search space
+        '''
+
+        # Reshape to use within a convolutional neural net.
+        # Last dimension is for "features" - there is only one here, since images are
+        # grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.
+        with tf.name_scope('reshape'):
+            try:
+                input_dim = int(math.sqrt(self.x_dim))
+            except:
+                print(
+                    'input dim cannot be sqrt and reshape. input dim: ' + str(self.x_dim))
+                logger.debug(
+                    'input dim cannot be sqrt and reshape. input dim: %s', str(self.x_dim))
+                raise
+            x_image = tf.reshape(self.images, [-1, input_dim, input_dim, 1])
+
+        """@nni.mutable_layers(
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=1, out_ch=self.channel_1_num),
+                               op.conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=7, in_ch=1, out_ch=self.channel_1_num),
+                               op.dilated_conv(in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=5, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=7, in_ch=1, out_ch=self.channel_1_num)],
+                fixed_inputs: [x_image],
+                layer_output: conv1_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_1_num)],
+                fixed_inputs: [conv1_out],
+                layer_output: post1_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post1_out],
+                layer_output: pool1_out
+            },
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.dilated_conv(in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=5, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num)],
+                fixed_inputs: [pool1_out],
+                optional_inputs: [post1_out],
+                optional_input_size: [0, 1],
+                layer_output: conv2_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_2_num)],
+                fixed_inputs: [conv2_out],
+                layer_output: post2_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post2_out],
+                optional_inputs: [post1_out, pool1_out],
+                optional_input_size: [0, 1],
+                layer_output: pool2_out
+            }
+        )"""
+
+        # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image
+        # is down to 7x7x64 feature maps -- maps this to 1024 features.
+        last_dim_list = pool2_out.get_shape().as_list()
+        assert(last_dim_list[1] == last_dim_list[2])
+        last_dim = last_dim_list[1]
+        with tf.name_scope('fc1'):
+            w_fc1 = op.weight_variable(
+                [last_dim * last_dim * self.channel_2_num, self.hidden_size])
+            b_fc1 = op.bias_variable([self.hidden_size])
+
+        h_pool2_flat = tf.reshape(
+            pool2_out, [-1, last_dim * last_dim * self.channel_2_num])
+        h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)
+
+        # Dropout - controls the complexity of the model, prevents co-adaptation of features.
+        with tf.name_scope('dropout'):
+            h_fc1_drop = tf.nn.dropout(h_fc1, self.keep_prob)
+
+        # Map the 1024 features to 10 classes, one for each digit
+        with tf.name_scope('fc2'):
+            w_fc2 = op.weight_variable([self.hidden_size, self.y_dim])
+            b_fc2 = op.bias_variable([self.y_dim])
+            y_conv = tf.matmul(h_fc1_drop, w_fc2) + b_fc2
+
+        with tf.name_scope('loss'):
+            cross_entropy = tf.reduce_mean(
+                tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=y_conv))
+        with tf.name_scope('adam_optimizer'):
+            self.train_step = tf.train.AdamOptimizer(
+                self.learning_rate).minimize(cross_entropy)
+
+        with tf.name_scope('accuracy'):
+            correct_prediction = tf.equal(
+                tf.argmax(y_conv, 1), tf.argmax(self.labels, 1))
+            self.accuracy = tf.reduce_mean(
+                tf.cast(correct_prediction, tf.float32))
+
+
+def download_mnist_retry(data_dir, max_num_retries=20):
+    """Try to download mnist dataset and avoid errors"""
+    for _ in range(max_num_retries):
+        try:
+            return input_data.read_data_sets(data_dir, one_hot=True)
+        except tf.errors.AlreadyExistsError:
+            time.sleep(1)
+    raise Exception("Failed to download MNIST.")
+
+def main(params):
+    '''
+    Main function, build mnist network, run and send result to NNI.
+    '''
+    # Import data
+    mnist = download_mnist_retry(params['data_dir'])
+    print('Mnist download data done.')
+    logger.debug('Mnist download data done.')
+
+    # Create the model
+    # Build the graph for the deep net
+    mnist_network = MnistNetwork(channel_1_num=params['channel_1_num'],
+                                 channel_2_num=params['channel_2_num'],
+                                 conv_size=params['conv_size'],
+                                 hidden_size=params['hidden_size'],
+                                 pool_size=params['pool_size'],
+                                 learning_rate=params['learning_rate'])
+    mnist_network.build_network()
+    logger.debug('Mnist build network done.')
+
+    # Write log
+    graph_location = tempfile.mkdtemp()
+    logger.debug('Saving graph to: %s', graph_location)
+    train_writer = tf.summary.FileWriter(graph_location)
+    train_writer.add_graph(tf.get_default_graph())
+
+    test_acc = 0.0
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        for i in range(params['batch_num']):
+            """@nni.training_update(tf, sess)"""
+            batch = mnist.train.next_batch(params['batch_size'])
+            mnist_network.train_step.run(feed_dict={mnist_network.images: batch[0],
+                                                    mnist_network.labels: batch[1],
+                                                    mnist_network.keep_prob: 1 - params['dropout_rate']}
+                                        )
+
+            if i % 100 == 0:
+                test_acc = mnist_network.accuracy.eval(
+                    feed_dict={mnist_network.images: mnist.test.images,
+                               mnist_network.labels: mnist.test.labels,
+                               mnist_network.keep_prob: 1.0})
+
+                """@nni.report_intermediate_result(test_acc)"""
+                logger.debug('test accuracy %g', test_acc)
+                logger.debug('Pipe send intermediate result done.')
+
+        test_acc = mnist_network.accuracy.eval(
+            feed_dict={mnist_network.images: mnist.test.images,
+                       mnist_network.labels: mnist.test.labels,
+                       mnist_network.keep_prob: 1.0})
+
+        """@nni.report_final_result(test_acc)"""
+        logger.debug('Final result is %g', test_acc)
+        logger.debug('Send final result done.')
+
+def get_params():
+    ''' Get parameters from command line '''
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_dir", type=str, default='/tmp/tensorflow/mnist/input_data', help="data directory")
+    parser.add_argument("--dropout_rate", type=float, default=0.5, help="dropout rate")
+    parser.add_argument("--channel_1_num", type=int, default=32)
+    parser.add_argument("--channel_2_num", type=int, default=64)
+    parser.add_argument("--conv_size", type=int, default=5)
+    parser.add_argument("--pool_size", type=int, default=2)
+    parser.add_argument("--hidden_size", type=int, default=1024)
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--batch_num", type=int, default=2000)
+    parser.add_argument("--batch_size", type=int, default=32)
+
+    args, _ = parser.parse_known_args()
+    return args
+
+if __name__ == '__main__':
+    try:
+        params = vars(get_params())
+        main(params)
+    except Exception as exception:
+        logger.exception(exception)
+        raise

examples/trials/mnist-nas/enas_mode/operators.py

+import tensorflow as tf
+import math
+
+
+def weight_variable(shape):
+    """weight_variable generates a weight variable of a given shape."""
+    initial = tf.truncated_normal(shape, stddev=0.1)
+    return tf.Variable(initial)
+
+def bias_variable(shape):
+    """bias_variable generates a bias variable of a given shape."""
+    initial = tf.constant(0.1, shape=shape)
+    return tf.Variable(initial)
+
+def sum_op(inputs):
+    """sum_op"""
+    fixed_input = inputs[0][0]
+    optional_input = tf.reduce_sum(inputs[1], axis=0)
+    if len(optional_input.get_shape()) < 1:
+        return fixed_input
+    fixed_shape = fixed_input.get_shape().as_list()
+    optional_shape = optional_input.get_shape().as_list()
+    assert fixed_shape[1] == fixed_shape[2]
+    assert optional_shape[1] == optional_shape[2]
+    pool_size = math.ceil(optional_shape[1] / fixed_shape[1])
+    pool_out = tf.nn.avg_pool(optional_input, ksize=[1, pool_size, pool_size, 1], strides=[1, pool_size, pool_size, 1], padding='SAME')
+    conv_matrix = weight_variable([1, 1, optional_shape[3], fixed_shape[3]])
+    conv_out = tf.nn.conv2d(pool_out, conv_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    return fixed_input + conv_out
+
+
+def conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """conv2d returns a 2d convolution layer with full stride."""
+    x_input = sum_op(inputs)
+    if size in [1, 3]:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.conv2d(x_input, w_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def twice_conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """twice_conv2d"""
+    x_input = sum_op(inputs)
+    if size in  [3, 7]:
+        w_matrix1 = weight_variable([1, size, in_ch, int(out_ch/2)])
+        out = tf.nn.conv2d(x_input, w_matrix1, strides=[1, 1, 1, 1], padding='SAME')
+        w_matrix2 = weight_variable([size, 1, int(out_ch/2), out_ch])
+        return tf.nn.conv2d(out, w_matrix2, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def dilated_conv(inputs, size=3, in_ch=-1, out_ch=-1):
+    """dilated_conv"""
+    x_input = sum_op(inputs)
+    if size == 3:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.atrous_conv2d(x_input, w_matrix, rate=2, padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def separable_conv(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """separable_conv"""
+    x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        depth_matrix = weight_variable([size, size, in_ch, 1])
+        point_matrix = weight_variable([1, 1, 1*in_ch, out_ch])
+        return tf.nn.separable_conv2d(x_input, depth_matrix, point_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def avg_pool(inputs, size=-1):
+    """avg_pool downsamples a feature map."""
+    x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.avg_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def max_pool(inputs, size=-1):
+    """max_pool downsamples a feature map."""
+    x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.max_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def post_process(inputs, ch_size=-1):
+    """post_process"""
+    x_input = inputs[0][0]
+    bias_matrix = bias_variable([ch_size])
+    return tf.nn.relu(x_input + bias_matrix)

examples/trials/mnist-nas/oneshot_mode/config_oneshot.yml

+authorName: default
+experimentName: example_mnist
+trialConcurrency: 1
+maxExecDuration: 1h
+maxTrialNum: 10
+#choice: local, remote, pai
+trainingServicePlatform: local
+#choice: true, false
+useAnnotation: true
+tuner:
+  codeDir: ../../../tuners/random_nas_tuner
+  classFileName: random_nas_tuner.py
+  className: RandomNASTuner
+trial:
+  command: python3 mnist-oneshot.py
+  codeDir: .
+  gpuNum: 0
+  nasMode: oneshot_mode

examples/trials/mnist-nas/oneshot_mode/mnist-oneshot.py

+"""A deep MNIST classifier using convolutional layers."""
+
+import argparse
+import logging
+import math
+import tempfile
+import time
+
+import tensorflow as tf
+from tensorflow.examples.tutorials.mnist import input_data
+
+import operators as op
+
+FLAGS = None
+
+logger = logging.getLogger('mnist_AutoML')
+
+
+class MnistNetwork(object):
+    '''
+    MnistNetwork is for initializing and building basic network for mnist.
+    '''
+    def __init__(self,
+                 channel_1_num,
+                 channel_2_num,
+                 conv_size,
+                 hidden_size,
+                 pool_size,
+                 learning_rate,
+                 x_dim=784,
+                 y_dim=10):
+        self.channel_1_num = channel_1_num
+        self.channel_2_num = channel_2_num
+        self.conv_size = conv_size
+        self.hidden_size = hidden_size
+        self.pool_size = pool_size
+        self.learning_rate = learning_rate
+        self.x_dim = x_dim
+        self.y_dim = y_dim
+
+        self.images = tf.placeholder(tf.float32, [None, self.x_dim], name='input_x')
+        self.labels = tf.placeholder(tf.float32, [None, self.y_dim], name='input_y')
+        self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')
+
+        self.train_step = None
+        self.accuracy = None
+
+    def build_network(self):
+        '''
+        Building network for mnist, meanwhile specifying its neural architecture search space
+        '''
+
+        # Reshape to use within a convolutional neural net.
+        # Last dimension is for "features" - there is only one here, since images are
+        # grayscale -- it would be 3 for an RGB image, 4 for RGBA, etc.
+        with tf.name_scope('reshape'):
+            try:
+                input_dim = int(math.sqrt(self.x_dim))
+            except:
+                print(
+                    'input dim cannot be sqrt and reshape. input dim: ' + str(self.x_dim))
+                logger.debug(
+                    'input dim cannot be sqrt and reshape. input dim: %s', str(self.x_dim))
+                raise
+            x_image = tf.reshape(self.images, [-1, input_dim, input_dim, 1])
+
+        """@nni.mutable_layers(
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=1, out_ch=self.channel_1_num),
+                               op.conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.twice_conv2d(size=7, in_ch=1, out_ch=self.channel_1_num),
+                               op.dilated_conv(in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=3, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=5, in_ch=1, out_ch=self.channel_1_num),
+                               op.separable_conv(size=7, in_ch=1, out_ch=self.channel_1_num)],
+                fixed_inputs: [x_image],
+                layer_output: conv1_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_1_num)],
+                fixed_inputs: [conv1_out],
+                layer_output: post1_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post1_out],
+                layer_output: pool1_out
+            },
+            {
+                layer_choice: [op.conv2d(size=1, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.twice_conv2d(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.dilated_conv(in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=3, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=5, in_ch=self.channel_1_num, out_ch=self.channel_2_num),
+                               op.separable_conv(size=7, in_ch=self.channel_1_num, out_ch=self.channel_2_num)],
+                fixed_inputs: [pool1_out],
+                optional_inputs: [post1_out],
+                optional_input_size: [0, 1],
+                layer_output: conv2_out
+            },
+            {
+                layer_choice: [op.post_process(ch_size=self.channel_2_num)],
+                fixed_inputs: [conv2_out],
+                layer_output: post2_out
+            },
+            {
+                layer_choice: [op.max_pool(size=3),
+                               op.max_pool(size=5),
+                               op.max_pool(size=7),
+                               op.avg_pool(size=3),
+                               op.avg_pool(size=5),
+                               op.avg_pool(size=7)],
+                fixed_inputs: [post2_out],
+                optional_inputs: [post1_out, pool1_out],
+                optional_input_size: [0, 1],
+                layer_output: pool2_out
+            }
+        )"""
+
+        # Fully connected layer 1 -- after 2 round of downsampling, our 28x28 image
+        # is down to 7x7x64 feature maps -- maps this to 1024 features.
+        last_dim_list = pool2_out.get_shape().as_list()
+        assert(last_dim_list[1] == last_dim_list[2])
+        last_dim = last_dim_list[1]
+        with tf.name_scope('fc1'):
+            w_fc1 = op.weight_variable(
+                [last_dim * last_dim * self.channel_2_num, self.hidden_size])
+            b_fc1 = op.bias_variable([self.hidden_size])
+
+        h_pool2_flat = tf.reshape(
+            pool2_out, [-1, last_dim * last_dim * self.channel_2_num])
+        h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, w_fc1) + b_fc1)
+
+        # Dropout - controls the complexity of the model, prevents co-adaptation of features.
+        with tf.name_scope('dropout'):
+            h_fc1_drop = tf.nn.dropout(h_fc1, self.keep_prob)
+
+        # Map the 1024 features to 10 classes, one for each digit
+        with tf.name_scope('fc2'):
+            w_fc2 = op.weight_variable([self.hidden_size, self.y_dim])
+            b_fc2 = op.bias_variable([self.y_dim])
+            y_conv = tf.matmul(h_fc1_drop, w_fc2) + b_fc2
+
+        with tf.name_scope('loss'):
+            cross_entropy = tf.reduce_mean(
+                tf.nn.softmax_cross_entropy_with_logits(labels=self.labels, logits=y_conv))
+        with tf.name_scope('adam_optimizer'):
+            self.train_step = tf.train.AdamOptimizer(
+                self.learning_rate).minimize(cross_entropy)
+
+        with tf.name_scope('accuracy'):
+            correct_prediction = tf.equal(
+                tf.argmax(y_conv, 1), tf.argmax(self.labels, 1))
+            self.accuracy = tf.reduce_mean(
+                tf.cast(correct_prediction, tf.float32))
+
+
+def download_mnist_retry(data_dir, max_num_retries=20):
+    """Try to download mnist dataset and avoid errors"""
+    for _ in range(max_num_retries):
+        try:
+            return input_data.read_data_sets(data_dir, one_hot=True)
+        except tf.errors.AlreadyExistsError:
+            time.sleep(1)
+    raise Exception("Failed to download MNIST.")
+
+def main(params):
+    '''
+    Main function, build mnist network, run and send result to NNI.
+    '''
+    # Import data
+    mnist = download_mnist_retry(params['data_dir'])
+    print('Mnist download data done.')
+    logger.debug('Mnist download data done.')
+
+    # Create the model
+    # Build the graph for the deep net
+    mnist_network = MnistNetwork(channel_1_num=params['channel_1_num'],
+                                 channel_2_num=params['channel_2_num'],
+                                 conv_size=params['conv_size'],
+                                 hidden_size=params['hidden_size'],
+                                 pool_size=params['pool_size'],
+                                 learning_rate=params['learning_rate'])
+    mnist_network.build_network()
+    logger.debug('Mnist build network done.')
+
+    # Write log
+    graph_location = tempfile.mkdtemp()
+    logger.debug('Saving graph to: %s', graph_location)
+    train_writer = tf.summary.FileWriter(graph_location)
+    train_writer.add_graph(tf.get_default_graph())
+
+    test_acc = 0.0
+    with tf.Session() as sess:
+        sess.run(tf.global_variables_initializer())
+        for i in range(params['batch_num']):
+            batch = mnist.train.next_batch(params['batch_size'])
+            mnist_network.train_step.run(feed_dict={mnist_network.images: batch[0],
+                                                    mnist_network.labels: batch[1],
+                                                    mnist_network.keep_prob: 1 - params['dropout_rate']}
+                                        )
+
+            if i % 100 == 0:
+                test_acc = mnist_network.accuracy.eval(
+                    feed_dict={mnist_network.images: mnist.test.images,
+                               mnist_network.labels: mnist.test.labels,
+                               mnist_network.keep_prob: 1.0})
+
+                """@nni.report_intermediate_result(test_acc)"""
+                logger.debug('test accuracy %g', test_acc)
+                logger.debug('Pipe send intermediate result done.')
+
+        test_acc = mnist_network.accuracy.eval(
+            feed_dict={mnist_network.images: mnist.test.images,
+                       mnist_network.labels: mnist.test.labels,
+                       mnist_network.keep_prob: 1.0})
+
+        """@nni.report_final_result(test_acc)"""
+        logger.debug('Final result is %g', test_acc)
+        logger.debug('Send final result done.')
+
+def get_params():
+    ''' Get parameters from command line '''
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_dir", type=str, default='/tmp/tensorflow/mnist/input_data', help="data directory")
+    parser.add_argument("--dropout_rate", type=float, default=0.5, help="dropout rate")
+    parser.add_argument("--channel_1_num", type=int, default=32)
+    parser.add_argument("--channel_2_num", type=int, default=64)
+    parser.add_argument("--conv_size", type=int, default=5)
+    parser.add_argument("--pool_size", type=int, default=2)
+    parser.add_argument("--hidden_size", type=int, default=1024)
+    parser.add_argument("--learning_rate", type=float, default=1e-4)
+    parser.add_argument("--batch_num", type=int, default=2000)
+    parser.add_argument("--batch_size", type=int, default=32)
+
+    args, _ = parser.parse_known_args()
+    return args
+
+if __name__ == '__main__':
+    try:
+        params = vars(get_params())
+        main(params)
+    except Exception as exception:
+        logger.exception(exception)
+        raise

examples/trials/mnist-nas/oneshot_mode/operators.py

+import tensorflow as tf
+import math
+
+
+def weight_variable(shape):
+    """weight_variable generates a weight variable of a given shape."""
+    initial = tf.truncated_normal(shape, stddev=0.1)
+    return tf.Variable(initial)
+
+def bias_variable(shape):
+    """bias_variable generates a bias variable of a given shape."""
+    initial = tf.constant(0.1, shape=shape)
+    return tf.Variable(initial)
+
+def sum_op(inputs):
+    """sum_op"""
+    fixed_input = inputs[0][0]
+    optional_input = tf.concat(inputs[1], axis=3)
+    fixed_shape = fixed_input.get_shape().as_list()
+    optional_shape = optional_input.get_shape().as_list()
+    assert fixed_shape[1] == fixed_shape[2]
+    assert optional_shape[1] == optional_shape[2]
+    pool_size = math.ceil(optional_shape[1] / fixed_shape[1])
+    pool_out = tf.nn.avg_pool(optional_input, ksize=[1, pool_size, pool_size, 1], strides=[1, pool_size, pool_size, 1], padding='SAME')
+    conv_matrix = weight_variable([1, 1, optional_shape[3], fixed_shape[3]])
+    conv_out = tf.nn.conv2d(pool_out, conv_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    return fixed_input + conv_out
+
+
+def conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """conv2d returns a 2d convolution layer with full stride."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [1, 3]:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.conv2d(x_input, w_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def twice_conv2d(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """twice_conv2d"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in  [3, 7]:
+        w_matrix1 = weight_variable([1, size, in_ch, int(out_ch/2)])
+        out = tf.nn.conv2d(x_input, w_matrix1, strides=[1, 1, 1, 1], padding='SAME')
+        w_matrix2 = weight_variable([size, 1, int(out_ch/2), out_ch])
+        return tf.nn.conv2d(out, w_matrix2, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def dilated_conv(inputs, size=3, in_ch=-1, out_ch=-1):
+    """dilated_conv"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size == 3:
+        w_matrix = weight_variable([size, size, in_ch, out_ch])
+        return tf.nn.atrous_conv2d(x_input, w_matrix, rate=2, padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def separable_conv(inputs, size=-1, in_ch=-1, out_ch=-1):
+    """separable_conv"""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        depth_matrix = weight_variable([size, size, in_ch, 1])
+        point_matrix = weight_variable([1, 1, 1*in_ch, out_ch])
+        return tf.nn.separable_conv2d(x_input, depth_matrix, point_matrix, strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def avg_pool(inputs, size=-1):
+    """avg_pool downsamples a feature map."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.avg_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+def max_pool(inputs, size=-1):
+    """max_pool downsamples a feature map."""
+    if not inputs[1]:
+        x_input = inputs[0][0]
+    else:
+        x_input = sum_op(inputs)
+    if size in [3, 5, 7]:
+        return tf.nn.max_pool(x_input, ksize=[1, size, size, 1], strides=[1, 1, 1, 1], padding='SAME')
+    else:
+        raise Exception("Unknown filter size: %d." % size)
+
+
+def post_process(inputs, ch_size=-1):
+    """post_process"""
+    x_input = inputs[0][0]
+    bias_matrix = bias_variable([ch_size])
+    return tf.nn.relu(x_input + bias_matrix)

src/nni_manager/rest_server/restValidationSchemas.ts

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

src/sdk/pynni/nni/__init__.py

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

src/sdk/pynni/nni/nas_utils.py

@@ -131,12 +131,41 @@ def oneshot_mode(
         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)
+    output_num = len(layer_outs)
+    rate = 0.01 ** (1 / output_num)
+    noise_shape = [output_num] + [1] * len(layer_outs[0].get_shape())
+    layer_outs = tf.nn.dropout(layer_outs, rate=rate, noise_shape=noise_shape)
+    layer_out = tf.reduce_sum(layer_outs, axis=0)
 
     return layer_out
 
 
-def reload_tensorflow_variables(session, tf=None):
+def darts_mode(
+        mutable_id,
+        mutable_layer_id,
+        funcs,
+        funcs_args,
+        fixed_inputs,
+        optional_inputs,
+        optional_input_size,
+        tf):
+    optional_inputs = list(optional_inputs.values())
+    layer_outs = [func([fixed_inputs, optional_inputs], **funcs_args[func_name])
+                  for func_name, func in funcs.items()]
+    # Create architecture weights for every func(op)
+    var_name = "{}_{}_".format(mutable_id, mutable_layer_id, "arch_weights")
+    if 'arch_logits_list' not in globals():
+        global arch_logits_list
+        arch_logits_list = list()
+    arch_logits = tf.get_variable(var_name, shape=[len[funcs]], trainable=False)
+    arch_logits_list.append(arch_logits)
+    arch_weights = tf.nn.softmax(arch_logits)
+    layer_out = tf.add_n([arch_weights[idx] * out for idx, out in enumerate(layer_outs)])
+
+    return layer_out
+
+
+def reload_tensorflow_variables(tf, session):
     '''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.
     ---------------
@@ -158,3 +187,22 @@ def reload_tensorflow_variables(session, tf=None):
             tf_variables[name_prefix]['funcs'].load(chosen_layer, session)
             tf_variables[name_prefix]['optional_inputs'].load(
                 chosen_inputs, session)
+
+
+def darts_training(tf, session, loss, feed_dict):
+    if 'optimizer' not in globals():
+        global arch_logits_list
+        global optimizer
+        global train_op
+        optimizer = tf.MomentumOptimizer(learning_rate=0.025)
+        # TODO: Calculate loss
+        grads_and_vars = optimizer.compute_gradients(loss, arch_logits_list)
+        train_op = optimizer.apply_gradients(grads_and_vars)
+    session.run(train_op)
+
+
+def training_update(nas_mode, tf=None, session=None, loss=None, feed_dict=None):
+    if nas_mode == 'darts_mode':
+        darts_training(tf, session, loss, feed_dict)
+    elif nas_mode == 'enas_mode':
+        reload_tensorflow_variables(tf, session)
\ No newline at end of file

src/sdk/pynni/nni/smartparam.py

@@ -24,7 +24,7 @@ import numpy as np
 
 from .env_vars import trial_env_vars
 from . import trial
-from .nas_utils import classic_mode, enas_mode, oneshot_mode
+from .nas_utils import classic_mode, enas_mode, oneshot_mode, darts_mode
 
 
 __all__ = [
@@ -158,8 +158,8 @@ else:
                             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'
+        assert tf is not None, 'Internal Error: Tensorflow should not be None in modes other than classic_mode'
+        if mode == 'enas_mode':
             return enas_mode(mutable_id,
                             mutable_layer_id,
                             funcs,
@@ -168,8 +168,7 @@ else:
                             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'
+        if mode == 'oneshot_mode':
             return oneshot_mode(mutable_id,
                             mutable_layer_id,
                             funcs,
@@ -178,7 +177,15 @@ else:
                             optional_inputs,
                             optional_input_size,
                             tf)
-        else:
+        if mode == 'darts_mode':
+            return darts_mode(mutable_id,
+                            mutable_layer_id,
+                            funcs,
+                            funcs_args,
+                            fixed_inputs,
+                            optional_inputs,
+                            optional_input_size,
+                            tf)
         raise RuntimeError('Unrecognized mode: %s' % mode)
 
     def _get_param(key):

test/config_test/examples/mnist-nas.test.yml

@@ -9,9 +9,10 @@ tuner:
   classFileName: random_nas_tuner.py
   className: RandomNASTuner
 trial:
-  codeDir: ../../../examples/trials/mnist-nas
+  codeDir: ../../../examples/trials/mnist-nas/classic_mode/
   command: python3 mnist.py --batch_num 100
   gpuNum: 0
+  nasMode: classic_mode
 
 useAnnotation: true
 multiPhase: false

tools/nni_annotation/code_generator.py

@@ -111,7 +111,7 @@ def parse_annotation_mutable_layers(code, lineno, nas_mode):
             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']:
+        if nas_mode in ['enas_mode', 'oneshot_mode', 'darts_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)
@@ -319,12 +319,11 @@ class Transformer(ast.NodeTransformer):
         else:
             return node  # not an annotation, ignore it
 
-        if string.startswith('@nni.get_next_parameter'):
-            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.training_update'):
+            expr = parse_annotation(string[1:])
+            call_node = expr.value
+            call_node.args.insert(0, ast.Str(s=self.nas_mode))
+            return expr
 
         if string.startswith('@nni.report_intermediate_result')  \
                 or string.startswith('@nni.report_final_result') \
@@ -378,8 +377,8 @@ def parse(code, nas_mode=None):
         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']:
+    # enas, oneshot and darts modes for tensorflow need tensorflow module, so we import it here
+    if nas_mode in ['enas_mode', 'oneshot_mode', 'darts_mode']:
         import_tf = ast.Import(names=[ast.alias(name='tensorflow', asname=None)])
         nodes.insert(last_future_import + 1, import_tf)
 

tools/nni_cmd/config_schema.py

@@ -221,7 +221,7 @@ common_trial_schema = {
     'command': setType('command', str),
     'codeDir': setPathCheck('codeDir'),
     Optional('gpuNum'): setNumberRange('gpuNum', int, 0, 99999),
-    Optional('nasMode'): setChoice('classic_mode', 'enas_mode', 'oneshot_mode')
+    Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode')
     }
 }
 
@@ -241,7 +241,7 @@ pai_trial_schema = {
     Optional('outputDir'): And(Regex(r'hdfs://(([0-9]{1,3}.){3}[0-9]{1,3})(:[0-9]{2,5})?(/.*)?'),\
                          error='ERROR: outputDir format error, outputDir format is hdfs://xxx.xxx.xxx.xxx:xxx'),
     Optional('virtualCluster'): setType('virtualCluster', str),
-    Optional('nasMode'): setChoice('classic_mode', 'enas_mode', 'oneshot_mode')
+    Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode')
     }
 }
 
@@ -256,7 +256,7 @@ pai_config_schema = {
 kubeflow_trial_schema = {
 'trial':{
         'codeDir':  setPathCheck('codeDir'),
-        Optional('nasMode'): setChoice('classic_mode', 'enas_mode', 'oneshot_mode'),
+        Optional('nasMode'): setChoice('nasMode', 'classic_mode', 'enas_mode', 'oneshot_mode', 'darts_mode'),
         Optional('ps'): {
             'replicas': setType('replicas', int),
             'command': setType('command', str),