def model_fn(features, labels, mode, params): ########### embedding ################# if not self.use_language_model: self.init_embedding() self.embed_query = self.embedding(features=features, name='x_query') else: self.embedding = None ############# encoder ################# #model params self.encoder.keep_prob = params['keep_prob'] self.encoder.is_training = params['is_training'] global_step = tf.train.get_or_create_global_step() if not self.use_language_model: out = self.encoder(self.embed_query, name='x_query', features=features) else: out = self.encoder(features=features) #pred = tf.nn.softmax(tf.layers.dense(out, self.num_class)) pred = tf.nn.softmax(out) ############### predict ################## if mode == tf.estimator.ModeKeys.PREDICT: predictions = { 'encode': out, 'logit': pred, 'label': features['label'] } return tf.estimator.EstimatorSpec(mode, predictions=predictions) ############### loss ################## loss = self.cal_loss(pred, labels, self.batch_size, self.conf) ############### train ################## if mode == tf.estimator.ModeKeys.TRAIN: if self.use_clr: self.learning_rate = cyclic_learning_rate( global_step=global_step, learning_rate=self.learning_rate, mode=self.clr_mode) optimizer = get_train_op(global_step, self.optimizer_type, loss, self.learning_rate, clip_grad=5) return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=optimizer) ############### eval ################## if mode == tf.estimator.ModeKeys.EVAL: eval_metric_ops = {} #{"accuracy": tf.metrics.accuracy( # labels=labels, predictions=predictions["classes"])} return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
def train_estimator_spec(self, mode, loss, global_step, params): """generate optimizer which can apply different learning rate in base model and downstream model,and return estimatorspec for training """ if self.use_clr: self.learning_rate = cyclic_learning_rate( global_step=global_step, learning_rate=self.learning_rate, max_lr=self.max_learning_rate, step_size=self.step_size, mode=self.clr_mode) else: self.learning_rate = tf.train.exponential_decay(self.learning_rate, global_step, 1, 0.99, staircase=True) optim_func = partial(get_train_op, global_step, self.optimizer_type, loss, clip_grad=5) if 'base_var' in params: #if contains base model variable list tvars = tf.trainable_variables() new_var_list = [] base_var_list = [] for var in tvars: name = var.name m = re.match("^(.*):\\d+$", name) if m is not None: name = m.group(1) if name in params['base_var']: base_var_list.append(var) continue new_var_list.append(var) optimizer_base = optim_func(learning_rate=self.base_learning_rate, var_list=base_var_list) optimizer_now = optim_func(learning_rate=self.learning_rate, var_list=new_var_list) if self.learning_rate == 0: raise ValueError('learning_rate can not be zero') if self.base_learning_rate == 0: # if base_learning_rate is set to be zero, than only # the downstream net parameters will be trained optimizer = optimizer_now else: optimizer = tf.group(optimizer_base, optimizer_now) else: optimizer = optim_func(learning_rate=self.learning_rate) return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=optimizer)
def model_fn(features, labels, mode, params): ############# embedding ################# if not self.use_language_model: self.embedding = init_embedding() if self.tfrecords_mode == 'class': self.embed_query = self.embedding(features=features, name='x_query') else: self.embed_query = self.embedding(features=features, name='x_query') self.embed_sample = self.embedding(features=features, name='x_sample') else: self.embedding = None ############# encoder ################# #model params self.encoder.keep_prob = params['keep_prob'] self.encoder.is_training = params['is_training'] global_step = tf.train.get_or_create_global_step() if self.sim_mode == 'cross': if not self.use_language_model: output = self.encoder(x_query=self.embed_query, x_sample=self.embed_sample, features=features) else: output = self.encoder(features=features) elif self.sim_mode == 'represent': if not self.use_language_model: #features['x_query_length'] = features['length'] output = self.encoder(self.embed_query, name='x_query', features=features) else: output = self.encoder(features=features) else: raise ValueError('unknown sim mode') ############### predict ################## if mode == tf.estimator.ModeKeys.PREDICT: #pdb.set_trace() predictions = { 'encode': output, 'pred': tf.cast(tf.greater(tf.nn.softmax(output)[:, 0], 0.5), tf.int32) if self.num_output == 2 else tf.cast(tf.greater(tf.nn.sigmoid(output), 0.5), tf.int32), 'score': tf.nn.softmax(output)[:, 0] if self.num_output == 2 else tf.nn.sigmoid(output), 'label': features['label'] } return tf.estimator.EstimatorSpec(mode, predictions=predictions) ############### loss ################## loss = cal_loss(output, labels, self.batch_size, self.conf) ############### train ################## if mode == tf.estimator.ModeKeys.TRAIN: if self.use_clr: self.learning_rate = cyclic_learning_rate( global_step=global_step, learning_rate=self.learning_rate, mode=self.clr_mode) optimizer = get_train_op(global_step, self.optimizer_type, loss, self.learning_rate, clip_grad=5) return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=optimizer) ############### eval ################## if mode == tf.estimator.ModeKeys.EVAL: eval_metric_ops = {} #{"accuracy": tf.metrics.accuracy( # labels=labels, predictions=predictions["classes"])} return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
def model_fn(features, labels, mode, params): if not self.use_language_model: self.init_embedding() if self.tfrecords_mode == 'class': self.embed_query = self.embedding(features = features, name = 'x_query') else: self.embed_query = self.embedding(features = features, name = 'x_query') self.embed_sample = self.embedding(features = features, name = 'x_sample') else: self.embedding = None #model params self.encoder.keep_prob = params['keep_prob'] self.encoder.is_training = params['is_training'] global_step = tf.train.get_or_create_global_step() if self.sim_mode == 'cross': if not self.use_language_model: pred = self.encoder(x_query = self.embed_query, x_sample = self.embed_sample, features = features) else: pred = self.encoder(features = features) elif self.sim_mode == 'represent': if not self.use_language_model: #features['x_query_length'] = features['length'] pred = self.encoder(self.embed_query, name = 'x_query', features = features) else: pred = self.encoder(features = features) else: raise ValueError('unknown sim mode') pos_target = tf.ones(shape = [int(self.batch_size/2)], dtype = tf.float32) neg_target = tf.zeros(shape = [int(self.batch_size/2)], dtype = tf.float32) if mode == tf.estimator.ModeKeys.PREDICT: predictions = { 'pred': pred, 'label': features['label'] } return tf.estimator.EstimatorSpec(mode, predictions=predictions) loss = self.cal_loss(pred, labels, pos_target, neg_target, self.batch_size, self.conf) if mode == tf.estimator.ModeKeys.TRAIN: if self.use_clr: self.learning_rate = cyclic_learning_rate(global_step=global_step, learning_rate = self.learning_rate, mode = self.clr_mode) optimizer = get_train_op(global_step, self.optimizer_type, loss, self.learning_rate, clip_grad = 5) return tf.estimator.EstimatorSpec(mode, loss = loss, train_op=optimizer) if mode == tf.estimator.ModeKeys.EVAL: eval_metric_ops = {} #{"accuracy": tf.metrics.accuracy( # labels=labels, predictions=predictions["classes"])} return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
def model_fn(features, labels, mode, params): ############# embedding ################# if not self.use_language_model: self.embedding = init_embedding() if self.tfrecords_mode == 'class': self.embed_query = self.embedding(features = features, name = 'x_query') else: self.embed_query = self.embedding(features = features, name = 'x_query') self.embed_sample = self.embedding(features = features, name = 'x_sample') else: self.embedding = None ############# encoder ################# #model params self.encoder.keep_prob = params['keep_prob'] self.encoder.is_training = params['is_training'] global_step = tf.train.get_or_create_global_step() if self.sim_mode == 'cross': if not self.use_language_model: output = self.encoder(x_query = self.embed_query, x_sample = self.embed_sample, features = features) else: output = self.encoder(features = features) elif self.sim_mode == 'represent': if not self.use_language_model: #features['x_query_length'] = features['length'] output = self.encoder(self.embed_query, name = 'x_query', features = features) else: output = self.encoder(features = features) else: raise ValueError('unknown sim mode') ############### predict ################## if mode == tf.estimator.ModeKeys.PREDICT: #pdb.set_trace() predictions = { 'encode': output, 'pred': tf.cast(tf.greater(tf.nn.softmax(output)[:,0], 0.5), tf.int32) if self.num_output == 2 else tf.cast(tf.greater(tf.nn.sigmoid(output), 0.5), tf.int32), 'score': tf.nn.softmax(output)[:,0] if self.num_output == 2 else tf.nn.sigmoid(output), 'label': features['label'] } return tf.estimator.EstimatorSpec(mode, predictions=predictions) ############### loss ################## loss = cal_loss(output, labels, self.batch_size, self.conf) ############### train ################## if mode == tf.estimator.ModeKeys.TRAIN: if self.use_clr: self.learning_rate = cyclic_learning_rate(global_step=global_step, learning_rate = self.learning_rate, mode = self.clr_mode) optim_func = partial(get_train_op, global_step, self.optimizer_type, loss, clip_grad =5) if 'base_var' in params: #if contains base model variable list tvars = tf.trainable_variables() now_var_list = [] sub_base_var_list = [] for var in tvars: name = var.name m = re.match("^(.*):\\d+$", name) if m is not None: name = m.group(1) if name in params['base_var']: sub_base_var_list.append(var) continue now_var_list.append(var) optimizer_base = optim_func(learning_rate = self.base_learning_rate, var_list = sub_base_var_list) optimizer_now = optim_func(learning_rate = self.learning_rate, var_list = now_var_list) optimizer = tf.group(optimizer_base, optimizer_now) else: optimizer = optim_func(learning_rate = self.learning_rate) return tf.estimator.EstimatorSpec(mode, loss = loss, train_op=optimizer) ############### eval ################## if mode == tf.estimator.ModeKeys.EVAL: eval_metric_ops = {} #{"accuracy": tf.metrics.accuracy( # labels=labels, predictions=predictions["classes"])} return tf.estimator.EstimatorSpec( mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)