def transfer_bert_model(bert_model_dir, output_bert_model):
graph = tf.Graph()
max_seq_len = 512
num_labels = 2
use_one_hot_embeddings = False
with graph.as_default():
with tf.Session() as sess:
input_ids = tf.placeholder(tf.int32, (None, None), 'input_ids')
input_mask = tf.placeholder(tf.int32, (None, None), 'input_mask')
segment_ids = tf.placeholder(tf.int32, (None, None), 'segment_ids')
bert_config = modeling.BertConfig.from_json_file(os.path.join(bert_model_dir, 'bert_config.json'))
model = modeling.BertModel(
config=bert_config,
is_training=False,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=use_one_hot_embeddings)
all_encoder_layers = model.get_all_encoder_layers()
input_x_bert_cls = model.get_pooled_output()
for idx, layer in enumerate(all_encoder_layers):
layer = tf.identity(layer, "encoder_layers_" + str(idx))
print("layer:", layer)
input_x_bert_cls = tf.identity(input_x_bert_cls, "input_x_bert_cls")
print("input_x_bert_cls", input_x_bert_cls)
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
saver.restore(sess, bert_model_dir + "/bert_model.ckpt")
saver.save(sess, output_bert_model)
def get_scaffold(self, mode, global_step=None, iter_initializer=None):
"""Get training scaffold."""
init_op = tf.global_variables_initializer()
if iter_initializer is None:
local_init_op = tf.tables_initializer()
else:
local_init_op = tf.group(tf.tables_initializer(), iter_initializer)
saver = self.get_saver(global_step)
scaffold = tf.train.Scaffold(
saver=saver, init_op=init_op, local_init_op=local_init_op)
return scaffold
def test_ngram_op_2_order(self):
''' test ngram 2-order op'''
ground_truth_2 = [0, 0, 0, 0, 0, 0, 0]
word_ngram = 2
t_input = tf.placeholder(shape=(4,), dtype=tf.int32)
t_ngram = py_x_ops.ngram(
t_input, word_ngrams=word_ngram, vocab_size=5000, bucket_size=100000)
logging.info("t_ngram: {}".format(t_ngram))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
ngram_result = sess.run(t_ngram, feed_dict={t_input: self.testcase[0]})
self.assertAllEqual(ngram_result, ground_truth_2)
def infer(self, **kwargs): # pylint: disable=arguments-differ, unused-argument
"""Make a inference."""
inputs = self.build_inputs(utils.INFER)
self.build()
self.session.run(tf.global_variables_initializer())
self.session.run(tf.tables_initializer())
self.session.run(inputs.iterator.initializer)
infer_data_size = self.config['data']['infer_data_size']
batch_size = self.config['data']['task']['batch_size']
steps = int(math.ceil(infer_data_size / batch_size))
weights_ckpt_dir = tf.train.latest_checkpoint(self.checkpoint_dir)
self.model.load_weights(weights_ckpt_dir)
logits = self.model.predict(inputs.input_x_dict, steps=steps)
preds = np.argmax(logits, axis=-1)
save_infer_res(self.config, logits, preds)
def export_model(self):
"""Export a model to tensorflow SavedModel."""
inputs = self.build_inputs(utils.INFER)
self.build()
logits = self.model(inputs.input_x_dict)
score = tf.nn.softmax(logits)
self.session.run(tf.global_variables_initializer())
self.session.run(tf.tables_initializer())
self.session.run(inputs.iterator.initializer)
weights_ckpt_dir = tf.train.latest_checkpoint(self.checkpoint_dir)
self.model.load_weights(weights_ckpt_dir)
output_dict = {"score": score}
to_saved_model(self.config, self.session, inputs.input_x_dict,
output_dict)
def eval(self):
"""Evaluate the model."""
inputs = self.build_inputs(utils.EVAL)
self.build()
self.session.run(tf.global_variables_initializer())
self.session.run(tf.tables_initializer())
self.session.run(inputs.iterator.initializer)
eval_data_size = self.config['data']['eval_data_size']
batch_size = self.config['data']['task']['batch_size']
steps = int(math.ceil(eval_data_size / batch_size))
weights_ckpt_dir = tf.train.latest_checkpoint(self.checkpoint_dir)
self.model.load_weights(weights_ckpt_dir)
results = self.model.evaluate(inputs.input_x_dict,
inputs.input_y_dict["input_y"],
steps=steps)
for metric, res in zip(self.model.metrics_names, results):
print("{}: {}".format(metric, res))
def test_batch_ngram_op_2_order(self):
''' tset batch 2-order ngram '''
ground_truth_2 = [[0, 0, 0, 0, 0, 0, 0], [223, 0, 0, 0, 0, 0, 0],
[0, 8, 5008, 0, 0, 0, 0], [4, 8, 102492, 0, 0, 0, 0],
[0, 0, 10, 5000, 5010, 0, 0],
[2, 5, 3, 103747, 51858, 0, 0],
[7, 2, 1, 24, 50599, 103743, 54395]]
word_ngram = 2
t_input = tf.placeholder(shape=(7, 4), dtype=tf.int32)
t_ngram = py_x_ops.ngram(
t_input, word_ngrams=word_ngram, vocab_size=5000, bucket_size=100000)
logging.info("batch t_ngram: {}".format(t_ngram))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
ngram_result = sess.run(t_ngram, feed_dict={t_input: self.testcase})
ngram_result = [list(res) for res in ngram_result]
self.assertAllEqual(ngram_result, ground_truth_2)
def train_core(self, train_inputs, eval_inputs=None):
"""Core part of training."""
self.build()
self.session.run(tf.global_variables_initializer())
self.session.run(tf.tables_initializer())
self.session.run(train_inputs.iterator.initializer)
if eval_inputs is not None:
self.session.run(eval_inputs.iterator.initializer)
validation_data = (eval_inputs.input_x_dict,
eval_inputs.input_y_dict["input_y"])
eval_data_size = self.config['data']['eval_data_size']
batch_size = self.config['data']['task']['batch_size']
validation_steps = int(eval_data_size / batch_size)
else:
validation_data = None
validation_steps = None
train_data_size = self.config['data']['train_data_size']
num_epochs = self.config['solver']['optimizer']['epochs']
batch_size = self.config['data']['task']['batch_size']
num_batch_per_epoch = int(math.ceil(train_data_size / batch_size))
callbacks = [
tf.keras.callbacks.TensorBoard(os.path.join(
self.model_path, "logs"),
histogram_freq=0,
write_graph=True,
write_grads=True,
write_images=True),
tf.keras.callbacks.ModelCheckpoint(os.path.join(
self.checkpoint_dir, "weights.{epoch:02d}"),
save_weights_only=True,
save_best_only=True)
]
self.model.fit(train_inputs.input_x_dict,
train_inputs.input_y_dict["input_y"],
callbacks=callbacks,
epochs=num_epochs,
steps_per_epoch=num_batch_per_epoch,
validation_data=validation_data,
validation_steps=validation_steps)
def main(_):
if FLAGS.checkpoints:
# Get the checkpoints list from flags and run some basic checks.
checkpoints = [c.strip() for c in FLAGS.checkpoints.split(",")]
checkpoints = [c for c in checkpoints if c]
if not checkpoints:
raise ValueError("No checkpoints provided for averaging.")
if FLAGS.prefix:
checkpoints = [FLAGS.prefix + c for c in checkpoints]
else:
assert FLAGS.num_last_checkpoints >= 1, "Must average at least one model"
assert FLAGS.prefix, ("Prefix must be provided when averaging last"
" N checkpoints")
checkpoint_state = tf.train.get_checkpoint_state(
os.path.dirname(FLAGS.prefix))
# Checkpoints are ordered from oldest to newest.
checkpoints = checkpoint_state.all_model_checkpoint_paths[
-FLAGS.num_last_checkpoints:]
checkpoints = [c for c in checkpoints if checkpoint_exists(c)]
if not checkpoints:
if FLAGS.checkpoints:
raise ValueError("None of the provided checkpoints exist. %s" %
FLAGS.checkpoints)
else:
raise ValueError("Could not find checkpoints at %s" %
os.path.dirname(FLAGS.prefix))
# Read variables from all checkpoints and average them.
logging.info("Reading variables and averaging checkpoints:")
for c in checkpoints:
logging.info("%s ", c)
var_list = tf.train.list_variables(checkpoints[0])
var_values, var_dtypes = {}, {}
for (name, shape) in var_list:
if not name.startswith("global_step"):
var_values[name] = np.zeros(shape)
for checkpoint in checkpoints:
reader = tf.train.load_checkpoint(checkpoint)
for name in var_values:
tensor = reader.get_tensor(name)
var_dtypes[name] = tensor.dtype
var_values[name] += tensor
logging.info("Read from checkpoint %s", checkpoint)
for name in var_values: # Average.
var_values[name] /= len(checkpoints)
with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
tf_vars = [
tf.get_variable(v, shape=var_values[v].shape, dtype=var_dtypes[v])
for v in var_values
]
placeholders = [tf.placeholder(v.dtype, shape=v.shape) for v in tf_vars]
assign_ops = [tf.assign(v, p) for (v, p) in zip(tf_vars, placeholders)]
global_step = tf.Variable(0,
name="global_step",
trainable=False,
dtype=tf.int64)
saver = tf.train.Saver(tf.all_variables())
# Build a model consisting only of variables, set them to the average values.
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
for p, assign_op, (name, value) in zip(placeholders, assign_ops,
six.iteritems(var_values)):
sess.run(assign_op, {p: value})
# Use the built saver to save the averaged checkpoint.
saver.save(sess, FLAGS.output_path, global_step=global_step)
logging.info("Averaged checkpoints saved in %s", FLAGS.output_path)