Suggestion for Neural Architecture Search with Reinforcement Learning

cmd/suggestion/nasrl/Dockerfile

@@ -1,7 +1,8 @@
-FROM python:3
+FROM python:3.6
 
 ADD . /usr/src/app/github.com/kubeflow/katib
 WORKDIR /usr/src/app/github.com/kubeflow/katib/cmd/suggestion/nasrl
+RUN pip install https://storage.googleapis.com/tensorflow/linux/cpu/tensorflow-1.12.0-cp36-cp36m-linux_x86_64.whl
 RUN pip install --no-cache-dir -r requirements.txt
 ENV PYTHONPATH /usr/src/app/github.com/kubeflow/katib:/usr/src/app/github.com/kubeflow/katib/pkg/api/python
 

cmd/suggestion/nasrl/requirements.txt

@@ -1,9 +1,3 @@
 grpcio
-duecredit
-cloudpickle==0.5.6
-numpy>=1.13.3
-scikit-learn>=0.19.0
-scipy>=0.19.1
-forestci
 protobuf
 googleapis-common-protos

pkg/suggestion/NAS_Reinforcement_Learning/Controller.py

@@ -0,0 +1,220 @@
+import tensorflow as tf
+from pkg.suggestion.NAS_Reinforcement_Learning.LSTM import stack_lstm
+from pkg.suggestion.NAS_Reinforcement_Learning.Trainer import get_train_ops
+
+
+class Controller(object):
+ def __init__(self,
+ num_layers=12,
+ num_operations=16,
+ lstm_size=64,
+ lstm_num_layers=1,
+ lstm_keep_prob=1.0,
+ tanh_constant=1.5,
+ temperature=None,
+ lr_init=1e-3,
+ lr_dec_start=0,
+ lr_dec_every=1000,
+ lr_dec_rate=0.9,
+ l2_reg=0,
+ entropy_weight=1e-4,
+ clip_mode=None,
+ grad_bound=None,
+ bl_dec=0.999,
+ optim_algo="adam",
+ sync_replicas=False,
+ num_aggregate=20,
+ num_replicas=1,
+ skip_target=0.4,
+ skip_weight=0.8,
+ name="controller"):
+
+ print("-" * 80)
+ print("Building Controller")
+
+ self.num_layers = num_layers
+ self.num_operations = num_operations
+
+ self.lstm_size = lstm_size
+ self.lstm_num_layers = lstm_num_layers
+ self.lstm_keep_prob = lstm_keep_prob
+ self.tanh_constant = tanh_constant
+ self.temperature = temperature
+ self.lr_init = lr_init
+ self.lr_dec_start = lr_dec_start
+ self.lr_dec_every = lr_dec_every
+ self.lr_dec_rate = lr_dec_rate
+ self.l2_reg = l2_reg
+ self.entropy_weight = entropy_weight
+ self.clip_mode = clip_mode
+ self.grad_bound = grad_bound
+ self.bl_dec = bl_dec
+
+ self.skip_target = skip_target
+ self.skip_weight = skip_weight
+
+ self.optim_algo = optim_algo
+ self.sync_replicas = sync_replicas
+ self.num_aggregate = num_aggregate
+ self.num_replicas = num_replicas
+ self.name = name
+
+ self._create_params()
+ self._build_sampler()
+
+ def _create_params(self):
+ initializer = tf.random_uniform_initializer(minval=-0.1, maxval=0.1)
+ with tf.variable_scope(self.name, initializer=initializer):
+ with tf.variable_scope("lstm"):
+ self.w_lstm = []
+ for layer_id in range(self.lstm_num_layers):
+ with tf.variable_scope("layer_{}".format(layer_id)):
+ w = tf.get_variable("w", [2 * self.lstm_size, 4 * self.lstm_size])
+ self.w_lstm.append(w)
+
+ self.g_emb = tf.get_variable("g_emb", [1, self.lstm_size])
+ with tf.variable_scope("emb"):
+ self.w_emb = tf.get_variable("w", [self.num_operations, self.lstm_size])
+ with tf.variable_scope("softmax"):
+ self.w_soft = tf.get_variable("w", [self.lstm_size, self.num_operations])
+
+ with tf.variable_scope("attention"):
+ self.w_attn_1 = tf.get_variable("w_1", [self.lstm_size, self.lstm_size])
+ self.w_attn_2 = tf.get_variable("w_2", [self.lstm_size, self.lstm_size])
+ self.v_attn = tf.get_variable("v", [self.lstm_size, 1])
+
+ def _build_sampler(self):
+ """Build the sampler ops and the log_prob ops."""
+
+ print("-" * 80)
+ print("Building Controller Sampler")
+ anchors = []
+ anchors_w_1 = []
+
+ arc_seq = []
+ entropys = []
+ log_probs = []
+ skip_count = []
+ skip_penaltys = []
+
+ prev_c = [tf.zeros([1, self.lstm_size], tf.float32) for _ in range(self.lstm_num_layers)]
+ prev_h = [tf.zeros([1, self.lstm_size], tf.float32) for _ in range(self.lstm_num_layers)]
+ inputs = self.g_emb
+ skip_targets = tf.constant([1.0 - self.skip_target, self.skip_target], dtype=tf.float32)
+ for layer_id in range(self.num_layers):
+ next_c, next_h = stack_lstm(inputs, prev_c, prev_h, self.w_lstm)
+ prev_c, prev_h = next_c, next_h
+ logit = tf.matmul(next_h[-1], self.w_soft)
+ if self.temperature is not None:
+ logit /= self.temperature
+ if self.tanh_constant is not None:
+ logit = self.tanh_constant * tf.tanh(logit)
+
+ operation_id = tf.multinomial(logit, 1)
+ operation_id = tf.to_int32(operation_id)
+ operation_id = tf.reshape(operation_id, [1])
+
+ arc_seq.append(operation_id)
+ log_prob = tf.nn.sparse_softmax_cross_entropy_with_logits(
+ logits=logit, labels=operation_id)
+ log_probs.append(log_prob)
+ entropy = tf.stop_gradient(log_prob * tf.exp(-log_prob))
+ entropys.append(entropy)
+ inputs = tf.nn.embedding_lookup(self.w_emb, operation_id)
+
+ next_c, next_h = stack_lstm(inputs, prev_c, prev_h, self.w_lstm)
+ prev_c, prev_h = next_c, next_h
+
+ if layer_id > 0:
+ query = tf.concat(anchors_w_1, axis=0)
+ query = tf.tanh(query + tf.matmul(next_h[-1], self.w_attn_2))
+ query = tf.matmul(query, self.v_attn)
+ logit = tf.concat([-query, query], axis=1)
+ if self.temperature is not None:
+ logit /= self.temperature
+ if self.tanh_constant is not None:
+ logit = self.tanh_constant * tf.tanh(logit)
+
+ skip = tf.multinomial(logit, 1)
+ skip = tf.to_int32(skip)
+ skip = tf.reshape(skip, [layer_id])
+ arc_seq.append(skip)
+
+ skip_prob = tf.sigmoid(logit)
+ kl = skip_prob * tf.log(skip_prob / skip_targets)
+ kl = tf.reduce_sum(kl)
+ skip_penaltys.append(kl)
+
+ log_prob = tf.nn.sparse_softmax_cross_entropy_with_logits(
+ logits=logit, labels=skip)
+ log_probs.append(tf.reduce_sum(log_prob, keepdims=True))
+
+ entropy = tf.stop_gradient(
+ tf.reduce_sum(log_prob * tf.exp(-log_prob), keepdims=True))
+ entropys.append(entropy)
+
+ skip = tf.to_float(skip)
+ skip = tf.reshape(skip, [1, layer_id])
+ skip_count.append(tf.reduce_sum(skip))
+ inputs = tf.matmul(skip, tf.concat(anchors, axis=0))
+ inputs /= (1.0 + tf.reduce_sum(skip))
+ else:
+ inputs = self.g_emb
+
+ anchors.append(next_h[-1])
+ anchors_w_1.append(tf.matmul(next_h[-1], self.w_attn_1))
+
+ arc_seq = tf.concat(arc_seq, axis=0)
+ self.sample_arc = tf.reshape(arc_seq, [-1])
+
+ entropys = tf.stack(entropys)
+ self.sample_entropy = tf.reduce_sum(entropys)
+
+ log_probs = tf.stack(log_probs)
+ self.sample_log_prob = tf.reduce_sum(log_probs)
+
+ skip_count = tf.stack(skip_count)
+ self.skip_count = tf.reduce_sum(skip_count)
+
+ skip_penaltys = tf.stack(skip_penaltys)
+ self.skip_penaltys = tf.reduce_mean(skip_penaltys)
+
+ def build_trainer(self):
+ self.reward = tf.placeholder(tf.float32, shape=())
+
+ normalize = tf.to_float(self.num_layers * (self.num_layers - 1) / 2)
+ self.skip_rate = tf.to_float(self.skip_count) / normalize
+
+ if self.entropy_weight is not None:
+ self.reward += self.entropy_weight * self.sample_entropy
+
+ self.sample_log_prob = tf.reduce_sum(self.sample_log_prob)
+ self.baseline = tf.Variable(0.0, dtype=tf.float32, trainable=False)
+ baseline_update = tf.assign_sub(self.baseline, (1 - self.bl_dec) * (self.baseline - self.reward))
+
+ with tf.control_dependencies([baseline_update]):
+ self.reward = tf.identity(self.reward)
+
+ self.loss = self.sample_log_prob * (self.reward - self.baseline)
+ if self.skip_weight is not None:
+ self.loss += self.skip_weight * self.skip_penaltys
+
+ self.train_step = tf.Variable(0, dtype=tf.int32, trainable=False, name=self.name + "_train_step")
+ tf_variables = [var for var in tf.trainable_variables() if var.name.startswith(self.name)]
+ print("-" * 80)
+
+ self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
+ self.loss,
+ tf_variables,
+ self.train_step,
+ clip_mode=self.clip_mode,
+ grad_bound=self.grad_bound,
+ l2_reg=self.l2_reg,
+ lr_init=self.lr_init,
+ lr_dec_start=self.lr_dec_start,
+ lr_dec_every=self.lr_dec_every,
+ lr_dec_rate=self.lr_dec_rate,
+ optim_algo=self.optim_algo,
+ sync_replicas=self.sync_replicas,
+ num_aggregate=self.num_aggregate,
+ num_replicas=self.num_replicas)

pkg/suggestion/NAS_Reinforcement_Learning/LSTM.py

@@ -0,0 +1,28 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import tensorflow as tf
+
+
+# TODO: will remove this function and use tf.nn.LSTMCell instead
+
+def lstm(x, prev_c, prev_h, w):
+ ifog = tf.matmul(tf.concat([x, prev_h], axis=1), w)
+ i, f, o, g = tf.split(ifog, 4, axis=1)
+ i = tf.sigmoid(i)
+ f = tf.sigmoid(f)
+ o = tf.sigmoid(o)
+ g = tf.tanh(g)
+ next_c = i * g + f * prev_c
+ next_h = o * tf.tanh(next_c)
+ return next_c, next_h
+
+
+def stack_lstm(x, prev_c, prev_h, w):
+ next_c, next_h = [], []
+ for layer_id, (_c, _h, _w) in enumerate(zip(prev_c, prev_h, w)):
+ inputs = x if layer_id == 0 else next_h[-1]
+ curr_c, curr_h = lstm(inputs, _c, _h, _w)
+ next_c.append(curr_c)
+ next_h.append(curr_h)
+ return next_c, next_h

pkg/suggestion/NAS_Reinforcement_Learning/Operation.py

@@ -0,0 +1,79 @@
+import itertools
+import numpy as np
+from pkg.api.python import api_pb2
+
+
+class Operation(object):
+ def __init__(self, opt_id, opt_type, opt_params):
+ self.opt_id = opt_id
+ self.opt_type = opt_type
+ self.opt_params = opt_params
+
+ def get_dict(self):
+ opt_dict = dict()
+ opt_dict['opt_id'] = self.opt_id
+ opt_dict['opt_type'] = self.opt_type
+ opt_dict['opt_params'] = self.opt_params
+ return opt_dict
+
+ def print_op(self, logger):
+ logger.info("Operation ID: \n\t{}".format(self.opt_id))
+ logger.info("Operation Type: \n\t{}".format(self.opt_type))
+ logger.info("Operations Parameters:")
+ for ikey in self.opt_params:
+ logger.info("\t{}: {}".format(ikey, self.opt_params[ikey]))
+ logger.info("")
+
+
+class SearchSpace(object):
+ def __init__(self, operations):
+ self.operation_list = list(operations.operation)
+ self.search_space = list()
+ self._parse_operations()
+ print()
+ self.num_operations = len(self.search_space)
+
+ def _parse_operations(self):
+ # search_sapce is a list of Operation class
+
+ operation_id = 0
+
+ for operation_dict in self.operation_list:
+ opt_type = operation_dict.operationType
+ opt_spec = list(operation_dict.parameter_configs.configs)
+ # avail_space is dict with the format {"spec_nam": [spec feasible values]}
+ avail_space = dict()
+ num_spec = len(opt_spec)
+
+ for ispec in opt_spec:
+ spec_name = ispec.name
+ if ispec.parameter_type == api_pb2.CATEGORICAL:
+ avail_space[spec_name] = list(ispec.feasible.list)
+ elif ispec.parameter_type == api_pb2.INT:
+ spec_min = int(ispec.feasible.min)
+ spec_max = int(ispec.feasible.max)
+ spec_step = int(ispec.feasible.step)
+ avail_space[spec_name] = range(spec_min, spec_max+1, spec_step)
+ elif ispec.parameter_type == api_pb2.DOUBLE:
+ spec_min = float(ispec.feasible.min)
+ spec_max = float(ispec.feasible.max)
+ spec_step = float(ispec.feasible.step)
+ if spec_step == 0:
+ print("Error, NAS Reinforcement Learning algorithm cannot accept continuous search space!")
+ exit(999)
+ double_list = np.arange(spec_min, spec_max+spec_step, spec_step)
+ if double_list[-1] > spec_max:
+ del double_list[-1]
+ avail_space[spec_name] = double_list
+
+ # generate all the combinations of possible operations
+ key_avail_space = list(avail_space.keys())
+ val_avail_space = list(avail_space.values())
+
+ for this_opt_vector in itertools.product(*val_avail_space):
+ opt_params = dict()
+ for i in range(num_spec):
+ opt_params[key_avail_space[i]] = this_opt_vector[i]
+ this_opt_class = Operation(operation_id, opt_type, opt_params)
+ self.search_space.append(this_opt_class)
+ operation_id += 1