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)
