def _ModelFn(features, labels, mode): if is_training: logits_out = self._BuildGraph(features) else: graph_def = self._GetGraphDef(use_trt, batch_size, model_dir) logits_out = importer.import_graph_def( graph_def, input_map={INPUT_NODE_NAME: features}, return_elements=[OUTPUT_NODE_NAME + ':0'], name='')[0] loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits_out) summary.scalar('loss', loss) classes_out = math_ops.argmax(logits_out, axis=1, name='classes_out') accuracy = metrics.accuracy(labels=labels, predictions=classes_out, name='acc_op') summary.scalar('accuracy', accuracy[1]) if mode == ModeKeys.EVAL: return EstimatorSpec(mode, loss=loss, eval_metric_ops={'accuracy': accuracy}) elif mode == ModeKeys.TRAIN: optimizer = AdamOptimizer(learning_rate=1e-2) train_op = optimizer.minimize(loss, global_step=get_global_step()) return EstimatorSpec(mode, loss=loss, train_op=train_op)
def model_fn(features, labels, mode, params): from tensorflow.python.estimator.model_fn import EstimatorSpec input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] label_ids = features["label_ids"] is_training = (mode == tf.estimator.ModeKeys.TRAIN) total_loss, logits, trans, pred_ids = self.create_model( bert_config, is_training, input_ids, input_mask, segment_ids, label_ids, num_labels, use_one_hot_embeddings) tvars = tf.trainable_variables() initialized_variable_names = {} if init_checkpoint: (assignment_map, initialized_variable_names) \ = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) if mode == tf.estimator.ModeKeys.TRAIN: #定义目标函数优化器 train_op = optimization.create_optimizer( total_loss, learning_rate, num_train_steps, num_warmup_steps, False) output_spec = EstimatorSpec(mode=mode, loss=total_loss, train_op=train_op) elif mode == tf.estimator.ModeKeys.EVAL: #验证集需要调整训练过程中的一下参数比如准确度 def metric_fn(per_example_loss, label_ids, logits): pass else: output_spec = EstimatorSpec(mode=mode, predictions=pred_ids) return output_spec
def model_fn(features, labels, mode, params): with tf.gfile.GFile(graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_names = ['input_ids', 'input_mask', 'input_type_ids'] output = tf.import_graph_def( graph_def, input_map={k + ':0': features[k] for k in input_names}, return_elements=['final_encodes:0']) if args.fp16: return EstimatorSpec(mode=mode, predictions={ 'unique_ids ': features['unique_ids'], 'encodes': tf.cast(output[0], tf.float32) }) else: return EstimatorSpec(mode=mode, predictions={ 'unique_ids ': features['unique_ids'], 'encodes': output[0] })
def ner_model_fn(features, labels, mode, params): """ 命名实体识别模型的model_fn :param features: :param labels: :param mode: :param params: :return: """ with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} #######将图从graph_def导入到当前默认图中,获取输出的pred_ids # input_map: 将graph_def中的输入名称(作为字符串) # 映射到张量对象的字典。输入图中指定的输入张量的值将被重新映射到相应的张量值。 # return_elements: 在graph_def中包含操作名的字符串列表,将作为operationobject返回; # 和 / 或graph_def中的张量名称,它们将作为张量对象返回 pred_ids = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_ids:0']) # pred_ids = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_ids']) # return EstimatorSpec(mode=mode, predictions={ # 'client_id': features['client_id'], # 'encodes': pred_ids[0] #TypeError: predictions[encodes] must be Tensor, given: name: "pred_ids" # }) print(pred_ids[0].shape.as_list(),' pred_id[0] shape') return EstimatorSpec(mode=mode, predictions={ 'client_id': features['client_id'], 'encodes': pred_ids[0] })
def model_fn(features, mode, params): from tensorflow.python.estimator.model_fn import EstimatorSpec input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] masked_lm_positions = features["masked_lm_positions"] masked_lm_ids = features["masked_lm_ids"] masked_lm_example_loss = self.create_model(bert_config, input_ids, input_mask, segment_ids, masked_lm_positions, masked_lm_ids, use_one_hot_embeddings) tvars = tf.trainable_variables() initialized_variable_names = {} if init_checkpoint: (assignment_map, initialized_variable_names) \ = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) output_spec = EstimatorSpec(mode=mode, predictions=masked_lm_example_loss) return output_spec
def train(features, labels, hparams, get_features_fn, embedding_fn, embedding_weights_initializer, bias_weights_initializer, global_rating_bias_initializer): """Constructs the matrix factorization model training graph.""" (query_movie_ids, query_movie_ratings, query_genre_ids, query_genre_freqs, query_genre_ratings, candidate_movie_id, candidate_genre_id) = (get_features_fn(features)) model_scores, _, _ = movie_candidate_score( query_movie_ids, query_movie_ratings, query_genre_ids, query_genre_freqs, query_genre_ratings, candidate_movie_id, candidate_genre_id, hparams, embedding_fn, embedding_weights_initializer, bias_weights_initializer, global_rating_bias_initializer) loss = tf.losses.mean_squared_error(features[LABEL_RATING_SCORE], model_scores) optimizer = tf.contrib.layers.OPTIMIZER_CLS_NAMES[hparams.optimizer]( learning_rate=hparams.learning_rate) if hparams.use_tpu: optimizer = tpu_optimizer.CrossShardOptimizer(optimizer) train_op = tf.contrib.layers.optimize_loss( loss=loss, summaries=[], global_step=tf.contrib.framework.get_global_step(), optimizer=optimizer, learning_rate=None) return EstimatorSpec(mode=TRAIN, predictions=model_scores, loss=loss, train_op=train_op)
def ner_model_fn(features, labels, mode, params): """ 命名实体识别模型的model_fn :param features: :param labels: :param mode: :param params: :return: """ with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} pred_ids = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_ids:0']) return EstimatorSpec(mode=mode, predictions={ 'client_id': features['client_id'], 'encodes': pred_ids[0] })
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" client_id = features["client_id"] input_ids = features["input_ids"] input_mask = features["input_mask"] input_type_ids = features["input_type_ids"] model = modeling.BertModel( config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, token_type_ids=input_type_ids, use_one_hot_embeddings=use_one_hot_embeddings) if mode != tf.estimator.ModeKeys.PREDICT: raise ValueError("Only PREDICT modes are supported: %s" % (mode)) tvars = tf.trainable_variables() (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) all_layers = [] if len(pooling_layer) == 1: encoder_layer = model.all_encoder_layers[pooling_layer[0]] else: for layer in pooling_layer: all_layers.append(model.all_encoder_layers[layer]) encoder_layer = tf.concat(all_layers, -1) input_mask = tf.cast(input_mask, tf.float32) if pooling_strategy == PoolingStrategy.REDUCE_MEAN: pooled = masked_reduce_mean(encoder_layer, input_mask) elif pooling_strategy == PoolingStrategy.REDUCE_MAX: pooled = masked_reduce_max(encoder_layer, input_mask) elif pooling_strategy == PoolingStrategy.REDUCE_MEAN_MAX: pooled = tf.concat([ masked_reduce_mean(encoder_layer, input_mask), masked_reduce_max(encoder_layer, input_mask) ], axis=1) elif pooling_strategy == PoolingStrategy.FIRST_TOKEN or pooling_strategy == PoolingStrategy.CLS_TOKEN: pooled = tf.squeeze(encoder_layer[:, 0:1, :], axis=1) elif pooling_strategy == PoolingStrategy.LAST_TOKEN or pooling_strategy == PoolingStrategy.SEP_TOKEN: seq_len = tf.cast(tf.reduce_sum(input_mask, axis=1), tf.int32) rng = tf.range(0, tf.shape(seq_len)[0]) indexes = tf.stack([rng, seq_len - 1], 1) pooled = tf.gather_nd(encoder_layer, indexes) elif pooling_strategy == PoolingStrategy.NONE: pooled = encoder_layer else: raise NotImplementedError() predictions = {'client_id': client_id, 'encodes': pooled} return EstimatorSpec(mode=mode, predictions=predictions)
def classification_model_fn(features, labels, mode, params): """ 文本分类模型的model_fn :param features: :param labels: :param mode: :param params: :return: """ with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} pred_probs = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_prob:0']) return EstimatorSpec( mode=mode, predictions={ 'client_id': features['client_id'], #'encodes': tf.argmax(pred_probs[0], axis=-1), #'score': tf.reduce_max(pred_probs[0], axis=-1) 'encodes': pred_probs[0], 'score': pred_probs[0] })
def ner_model_fn(features, labels, mode, params): """ 命名实体识别模型的model_fn :param features: :param labels: :param mode: :param params: :return: """ # with tf.gfile.GFile(self.graph_path, 'rb') as f: # graph_def = tf.GraphDef() # graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["input_type_ids"] print("++++++++++++++",input_ids) # tokens =features["tokens"] # print("++++++++++++++",tokens) input_map = {"input_ids": input_ids, "input_mask": input_mask,"segment_ids":segment_ids} predicate_probabilities, token_label_probabilities= tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['predicate_probabilities:0','token_label_probabilities:0']) predicate_index = tf.argmax(predicate_probabilities, axis=-1, output_type=tf.int32) token_label_index = tf.argmax(token_label_probabilities, axis=-1) return EstimatorSpec(mode=mode, predictions={ # 'client_id': features['client_id'], 'predicate_index': predicate_index, 'token_label_index': token_label_index, 'predicate_probabilities': predicate_probabilities, 'token_label_probabilities': token_label_probabilities, # "tokens":tokens })
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" unique_ids = features["unique_ids"] input_ids = features["input_ids"] input_mask = features["input_mask"] input_type_ids = features["input_type_ids"] model = modeling.BertModel( config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, token_type_ids=input_type_ids, use_one_hot_embeddings=use_one_hot_embeddings) if mode != tf.estimator.ModeKeys.PREDICT: raise ValueError("Only PREDICT modes are supported: %s" % (mode)) tvars = tf.trainable_variables() (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) tf.logging.info("**** Trainable Variables ****") for var in tvars: init_string = "" if var.name in initialized_variable_names: init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) #all_layers = model.get_all_encoder_layers() encoder_layer = model.all_encoder_layers[pooling_layer] if pooling_strategy == PoolingStrategy.REDUCE_MEAN: pooled = tf.reduce_mean(encoder_layer, axis=1) elif pooling_strategy == PoolingStrategy.REDUCE_MAX: pooled = tf.reduce_max(encoder_layer, axis=1) elif pooling_strategy == PoolingStrategy.REDUCE_MEAN_MAX: pooled = tf.concat([ tf.reduce_max(encoder_layer, axis=1), tf.reduce_max(encoder_layer, axis=1) ], axis=1) elif pooling_strategy == PoolingStrategy.FIRST_TOKEN or pooling_strategy == PoolingStrategy.CLS_TOKEN: pooled = tf.squeeze(encoder_layer[:, 0:1, :], axis=1) elif pooling_strategy == PoolingStrategy.LAST_TOKEN or pooling_strategy == PoolingStrategy.SEP_TOKEN: seq_len = tf.cast(tf.reduce_sum(input_mask, axis=1), tf.int32) rng = tf.range(0, tf.shape(seq_len)[0]) indexes = tf.stack([rng, seq_len - 1], 1) pooled = tf.gather_nd(encoder_layer, indexes) else: raise NotImplementedError() predictions = {"unique_id": unique_ids, "encodes": pooled} return EstimatorSpec(mode=mode, predictions=predictions)
def model_fn(features, mode): with tf.gfile.GFile(self._graphdef, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) output = tf.import_graph_def(graph_def, input_map={k + ':0': features[k] for k in self._input_names}, return_elements=['final_encodes:0']) return EstimatorSpec(mode=mode, predictions={'output': output[0]})
def model_fn(features, labels, mode, params): predicts = tf.import_graph_def( graph_def, input_map={k + ":0": features[k] for k in input_names}, return_elements=['output/predictions:0', 'output/scores:0']) return EstimatorSpec(mode=mode, predictions={ "predict": predicts[0], "scores": predicts[1] })
def classification_model_fn(self, features, mode): with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} pred_probs = tf.import_graph_def(graph_def, name='', input_map=input_map, return_elements=['pred_prob:0']) return EstimatorSpec(mode=mode, predictions={ 'encodes': tf.argmax(pred_probs[0], axis=-1), 'score': tf.reduce_max(pred_probs[0], axis=-1)})
def model_fn(features, labels, mode, params): del params x = core_layers.dense(features, 100) x = core_layers.dense(x, 100) x = core_layers.dense(x, 100) x = core_layers.dense(x, 100) y = core_layers.dense(x, 1) loss = losses.mean_squared_error(labels, y) opt = adam.AdamOptimizer(learning_rate=0.1) train_op = opt.minimize( loss, global_step=training_util.get_or_create_global_step()) return EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" tf.logging.info("*** Features ***") for name in sorted(features.keys()): tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) unique_ids = features["unique_ids"] input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] is_training = (mode == tf.estimator.ModeKeys.TRAIN) (start_logits, end_logits) = create_model( bert_config=bert_config, is_training=is_training, input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids, use_one_hot_embeddings=use_one_hot_embeddings) tvars = tf.trainable_variables() initialized_variable_names = {} scaffold_fn = None if init_checkpoint: (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) tf.logging.info("**** Trainable Variables ****") for var in tvars: init_string = "" if var.name in initialized_variable_names: init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) predictions = { "unique_ids": unique_ids, "start_logits": start_logits, "end_logits": end_logits, } output_spec = EstimatorSpec(mode=mode, predictions=predictions, scaffold=scaffold_fn) return output_spec
def model_fn(features, labels, mode, params): with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) # edit by gavin: 加载模型中的图 input_names = ['input_ids', 'input_mask', 'input_type_ids'] output = tf.import_graph_def(graph_def, input_map={k + ':0': features[k] for k in input_names}, return_elements=['final_encodes:0']) return EstimatorSpec(mode=mode, predictions={ 'encodes': output[0] })
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" unique_ids = features["unique_ids"] input_ids = features["input_ids"] input_mask = features["input_mask"] input_type_ids = features["input_type_ids"] jit_scope = tf.contrib.compiler.jit.experimental_jit_scope with jit_scope(): model = modeling.BertModel(config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, token_type_ids=input_type_ids) if mode != tf.estimator.ModeKeys.PREDICT: raise ValueError("Only PREDICT modes are supported: %s" % (mode)) tvars = tf.trainable_variables() (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) tf.logging.info("**** Trainable Variables ****") for var in tvars: init_string = "" if var.name in initialized_variable_names: init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) all_layers = model.get_all_encoder_layers() predictions = { "unique_id": unique_ids, } for (i, layer_index) in enumerate(layer_indexes): predictions["layer_output_%d" % i] = all_layers[layer_index] from tensorflow.python.estimator.model_fn import EstimatorSpec output_spec = EstimatorSpec(mode=mode, predictions=predictions) return output_spec
def model_fn(features, labels, mode, params): predicts = tf.import_graph_def( graph_def, input_map={k + ":0": features[k] for k in input_names}, return_elements=[ 'transformer-encoding/encoded_outputs:0', 'output/softmax_scores:0', 'output/predictions:0' ]) return EstimatorSpec(mode=mode, predictions={ "encoded_outputs": predicts[0], "scores": predicts[1], "predictions": predicts[2] })
def model_fn(features, labels, mode, params): with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} pred_theme, embeddings = tf.import_graph_def( graph_def, name='', input_map=input_map, return_elements=['theme/logits/ArgMax:0', 'bert/transformer/Reshape_2:0']) return EstimatorSpec(mode=mode, predictions={ 'client_id': features['client_id'], 'pred_themes': pred_theme, 'embeddings': embeddings })
def model_fn(features, labels, mode, params): with tf.gfile.GFile(graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_names = ['input_ids', 'input_mask', 'input_type_ids'] encoder_layer = tf.import_graph_def(graph_def, input_map={k + ':0': features[k] for k in input_names}, return_elements=['final_encodes:0']) predictions = { # 'client_id': client_id, 'encodes': encoder_layer[0] } return EstimatorSpec(mode=mode, predictions=predictions)
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument from tensorflow.python.estimator.model_fn import EstimatorSpec tf.logging.info("*** Features ***") for name in sorted(features.keys()): tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] label_ids = features["label_ids"] is_training = (mode == tf.estimator.ModeKeys.TRAIN) (total_loss, per_example_loss, logits, probabilities) = create_model(bert_config, is_training, input_ids, input_mask, segment_ids, label_ids, num_labels, use_one_hot_embeddings) tvars = tf.trainable_variables() initialized_variable_names = {} print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") if init_checkpoint: print("?????????????????????????????????????????????????????", init_checkpoint) (assignment_map, initialized_variable_names) \ = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) tf.logging.info("**** Trainable Variables ****") with tf.Session() as sess: sess.run(tf.global_variables_initializer()) for var in tvars: init_string = "" if var.name in initialized_variable_names: if len(var.shape) == 2: print("var.shape=2", var.name, sess.run(var)[0][1]) else: print("var.shape=1", var.name, sess.run(var)[1]) init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) output_spec = EstimatorSpec(mode=mode, predictions=probabilities) return output_spec
def model_fn(features, labels, mode, params): with tf.gfile.GFile(self.graph_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) # input_names = ['input_ids', 'input_mask', 'input_type_ids'] input_names = list(self.get_tf_feature_inputs(tf).keys()) output_names = [self.get_tf_feature_output(tf)] output = tf.import_graph_def(graph_def, input_map={k + ':0': features[k] for k in input_names}, return_elements=[k + ':0' for k in output_names]) return EstimatorSpec(mode=mode, predictions={ 'client_id': features['client_id'], 'outputs': output[0] })
def model_fn(features, labels, mode, params): with tf.gfile.GFile(self.graph_path, 'rb') as f: ###从 python Graph中序列化出来的图就叫做 GraphDef graph_def = tf.GraphDef() graph_def.ParseFromString(f.read())####读取pb格式的graph # input_names = ['input_ids', 'input_mask', 'input_type_ids']###我在训练和压缩的时候就没input_type_ids input_names = ['input_ids', 'input_mask'] ###f.import_graph_def():将图从graph_def导入到当前默认图中 output = tf.import_graph_def(graph_def, input_map={k + ':0': features[k] for k in input_names}, return_elements=['final_encodes:0']) return EstimatorSpec(mode=mode, predictions={ 'client_id': features['client_id'], 'encodes': output[0] })
def classification_model_fn( self, features, mode, ): input_ids = features["input_ids"] input_mask = features["input_mask"] input_map = {"input_ids": input_ids, "input_mask": input_mask} pred_probs = tf.import_graph_def(self.runfunc, name='', input_map=input_map, return_elements=['pred_prob:0']) return EstimatorSpec(mode=mode, predictions={ 'encodes': tf.argmax(pred_probs[0], axis=-1), 'score': tf.reduce_max(pred_probs[0], axis=-1) })
def model_fn(features, labels, mode, params): with tf.gfile.GFile(model_file, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) input_names = ['input_x'] predicts = tf.import_graph_def( graph_def, input_map={k + ":0": features[k] for k in input_names}, return_elements=[ 'transformer-encoding/encoded_outputs:0', 'output/softmax_scores:0', 'output/predictions:0' ]) return EstimatorSpec(mode=mode, predictions={ "encoded_outputs": predicts[0], "scores": predicts[1], "predictions": predicts[2] })
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" # unique_ids = features["unique_ids"] input_ids = features["input_ids"] input_mask = features["input_mask"] input_type_ids = features["input_type_ids"] model = modeling.BertModel( config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, token_type_ids=input_type_ids, use_one_hot_embeddings=use_one_hot_embeddings) if mode != tf.estimator.ModeKeys.PREDICT: raise ValueError("Only PREDICT modes are supported: %s" % (mode)) tvars = tf.trainable_variables() (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) # tf.logging.info("**** Trainable Variables ****") # for var in tvars: # init_string = "" # if var.name in initialized_variable_names: # init_string = ", *INIT_FROM_CKPT*" # tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, # init_string) # predictions = { # 'unique_id': unique_ids, # 'pooled': model.get_pooled_output() # } return EstimatorSpec(mode=mode, predictions=model.get_sentence_encoding())
def train(features, labels, hparams, embedding_weights_initializer, class_weights_initializer, class_biases_initializer, get_features_fn, embedding_fn): """Constructs the training graph.""" (movie_ids_ratings, genre_ids_freqs, genre_ids_ratings) = (get_features_fn(features)) query_embeddings = embed_query_features(movie_ids_ratings, genre_ids_freqs, genre_ids_ratings, hparams, TRAIN, embedding_weights_initializer, embedding_fn) class_weights, class_biases = class_weights_biases( hparams, class_weights_initializer, class_biases_initializer) scores = tf.matmul(query_embeddings, tf.transpose(class_weights)) + class_biases target_one_hot = tf.one_hot(indices=features['candidate_movie_id_values'], depth=MOVIE_VOCAB_SIZE, on_value=1.0) loss = tf.reduce_mean( tf.nn.softmax_cross_entropy_with_logits(labels=target_one_hot, logits=scores)) optimizer = tf.contrib.layers.OPTIMIZER_CLS_NAMES[hparams.optimizer]( learning_rate=hparams.learning_rate) if hparams.use_tpu: optimizer = tpu_optimizer.CrossShardOptimizer(optimizer) train_op = tf.contrib.layers.optimize_loss( loss=loss, summaries=[], global_step=tf.contrib.framework.get_global_step(), optimizer=optimizer, learning_rate=None) return EstimatorSpec(mode=TRAIN, predictions=scores, loss=loss, train_op=train_op)
def model_fn(features, labels, mode, params): from tensorflow.python.estimator.model_fn import EstimatorSpec input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] label_ids = features["label_ids"] is_training = (mode == tf.estimator.ModeKeys.TRAIN) probabilities = self.create_model(bert_config, input_ids, input_mask, segment_ids, label_ids, num_labels, use_one_hot_embeddings) tvars = tf.trainable_variables() if init_checkpoint: (assignment_map, initialized_variable_names) \ = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint) tf.train.init_from_checkpoint(init_checkpoint, assignment_map) output_spec = EstimatorSpec(mode=mode, predictions=probabilities) return output_spec
def _build(self, features, labels=None, params=None, config=None): # Pre-process features and labels features, labels = self._preprocess(features, labels) results = self._call_graph_fn(inputs=features) if not isinstance(results, BridgeSpec): raise ValueError('`bridge_fn` should return a BridgeSpec.') loss = None train_op = None eval_metrics = None if Modes.is_infer(self.mode): predictions = self._build_predictions(results=results.results, features=features, labels=labels) else: losses, loss = self._build_loss(results, features, features) eval_metrics = self._build_eval_metrics(results.results, features, features) if Modes.is_train(self.mode): train_op = self._build_train_op(loss) self._build_summary_op(results=results.results, features=features, labels=labels) predictions = self._build_predictions(results=results.results, features=features, labels=labels, losses=losses) # We add 'useful' tensors to the graph collection so that we # can easly find them in our hooks/monitors. track(predictions, tf.GraphKeys.PREDICTIONS) return EstimatorSpec(mode=self.mode, predictions=predictions, loss=loss, train_op=train_op, eval_metric_ops=eval_metrics)