def model_fn(features, labels, mode):
model = bert_encoder(model_config,
features,
labels,
mode,
target,
reuse=model_reuse)
label_ids = features["label_ids"]
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=model_reuse):
(loss, per_example_loss,
logits) = classifier.classifier(model_config,
model.get_pooled_output(),
num_labels, label_ids,
dropout_prob)
temperature_log_prob = tf.nn.log_softmax(
logits / kargs.get("temperature", 2))
return [loss, per_example_loss, logits, temperature_log_prob]
def model_fn(features, labels, mode):
if target:
input_ids = features["input_ids_{}".format(target)]
input_mask = features["input_mask_{}".format(target)]
segment_ids = features["segment_ids_{}".format(target)]
else:
input_ids = features["input_ids"]
input_mask = features["input_mask"]
segment_ids = features["segment_ids"]
input_shape = bert_utils.get_shape_list(input_ids, expected_rank=3)
batch_size = input_shape[0]
choice_num = input_shape[1]
seq_length = input_shape[2]
if target:
real_features = {
"input_ids_{}".format(target):
tf.reshape(input_ids, [batch_size * choice_num, seq_length]),
"input_mask_{}".format(target):
tf.reshape(input_mask, [batch_size * choice_num, seq_length]),
"segment_ids_{}".format(target):
tf.reshape(segment_ids, [batch_size * choice_num, seq_length]),
"label_ids":
features["label_ids"]
}
else:
real_features = {
"input_ids":
tf.reshape(input_ids, [batch_size * choice_num, seq_length]),
"input_mask":
tf.reshape(input_mask, [batch_size * choice_num, seq_length]),
"segment_ids":
tf.reshape(segment_ids, [batch_size * choice_num, seq_length]),
"label_ids":
features["label_ids"]
}
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
model = bert_encoder(model_config,
real_features,
labels,
mode,
target,
reuse=model_reuse)
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=reuse):
(loss, per_example_loss,
logits) = classifier.multi_choice_classifier(
model_config, model.get_pooled_output(), num_labels,
label_ids, dropout_prob)
tvars = model_io_fn.get_params(scope,
not_storage_params=not_storage_params)
if load_pretrained:
model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope)
model_io_fn.set_saver(var_lst=tvars)
if mode == tf.estimator.ModeKeys.TRAIN:
model_io_fn.print_params(tvars, string=", trainable params")
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer_fn.get_train_op(
loss, tvars, opt_config.init_lr,
opt_config.num_train_steps)
estimator_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
train_op=train_op)
return {
"estimator_spec": estimator_spec,
"train": {
"loss": loss,
"logits": logits,
"train_op": train_op
}
}
elif mode == tf.estimator.ModeKeys.PREDICT:
print(logits.get_shape(), "===logits shape===")
pred_label = tf.argmax(logits, axis=-1, output_type=tf.int32)
prob = tf.nn.softmax(logits)
max_prob = tf.reduce_max(prob, axis=-1)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
'pred_label': pred_label,
"max_prob": max_prob
},
export_outputs={
"output":
tf.estimator.export.PredictOutput({
'pred_label': pred_label,
"max_prob": max_prob
})
})
return {"estimator_spec": estimator_spec}
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss, logits, label_ids):
"""Computes the loss and accuracy of the model."""
sentence_log_probs = tf.reshape(logits, [-1, logits.shape[-1]])
sentence_predictions = tf.argmax(logits,
axis=-1,
output_type=tf.int32)
sentence_labels = tf.reshape(label_ids, [-1])
sentence_accuracy = tf.metrics.accuracy(
labels=label_ids, predictions=sentence_predictions)
sentence_mean_loss = tf.metrics.mean(values=per_example_loss)
sentence_f = tf_metrics.f1(label_ids,
sentence_predictions,
num_labels,
label_lst,
average="macro")
eval_metric_ops = {
"f1": sentence_f,
"loss": sentence_mean_loss,
"acc": sentence_accuracy
}
return eval_metric_ops
eval_metric_ops = metric_fn(per_example_loss, logits, label_ids)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
return {
"estimator_spec": estimator_spec,
"eval": {
"per_example_loss": per_example_loss,
"logits": logits
}
}
else:
raise NotImplementedError()
def model_fn(features, labels, mode):
shape_lst_a = bert_utils.get_shape_list(features['input_ids_a'])
batch_size_a = shape_lst_a[0]
total_length_a = shape_lst_a[1]
shape_lst_b = bert_utils.get_shape_list(features['input_ids_b'])
batch_size_b = shape_lst_b[0]
total_length_b = shape_lst_b[1]
features['input_ids_a'] = tf.reshape(features['input_ids_a'],
[-1, model_config.max_length])
features['segment_ids_a'] = tf.reshape(features['segment_ids_a'],
[-1, model_config.max_length])
features['input_mask_a'] = tf.cast(
tf.not_equal(features['input_ids_a'], kargs.get('[PAD]', 0)),
tf.int64)
features['input_ids_b'] = tf.reshape(
features['input_ids_b'],
[-1, model_config.max_predictions_per_seq])
features['segment_ids_b'] = tf.reshape(
features['segment_ids_b'],
[-1, model_config.max_predictions_per_seq])
features['input_mask_b'] = tf.cast(
tf.not_equal(features['input_ids_b'], kargs.get('[PAD]', 0)),
tf.int64)
features['batch_size'] = batch_size_a
features['total_length_a'] = total_length_a
features['total_length_b'] = total_length_b
model_dict = {}
for target in ["a", "b"]:
model = bert_encoder(model_config,
features,
labels,
mode,
target,
reuse=tf.AUTO_REUSE)
model_dict[target] = model
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=model_reuse):
(loss, per_example_loss, logits,
transition_params) = multi_position_crf_classifier(
model_config, features, model_dict, num_labels, dropout_prob)
model_io_fn = model_io.ModelIO(model_io_config)
tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
try:
params_size = model_io_fn.count_params(model_config.scope)
print("==total params==", params_size)
except:
print("==not count params==")
print(tvars)
if load_pretrained == "yes":
model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope)
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer_fn = optimizer.Optimizer(opt_config)
model_io_fn.print_params(tvars, string=", trainable params")
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
print("==update_ops==", update_ops)
with tf.control_dependencies(update_ops):
train_op = optimizer_fn.get_train_op(
loss, tvars, opt_config.init_lr,
opt_config.num_train_steps, **kargs)
train_op, hooks = model_io_fn.get_ema_hooks(
train_op,
tvars,
kargs.get('params_moving_average_decay', 0.99),
scope,
mode,
first_stage_steps=opt_config.num_warmup_steps,
two_stage=True)
model_io_fn.set_saver()
if kargs.get("task_index", 1) == 0 and kargs.get(
"run_config", None):
training_hooks = []
elif kargs.get("task_index", 1) == 0:
model_io_fn.get_hooks(kargs.get("checkpoint_dir", None),
kargs.get("num_storage_steps", 1000))
training_hooks = model_io_fn.checkpoint_hook
else:
training_hooks = []
if len(optimizer_fn.distributed_hooks) >= 1:
training_hooks.extend(optimizer_fn.distributed_hooks)
print(training_hooks, "==training_hooks==", "==task_index==",
kargs.get("task_index", 1))
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
print(tf.global_variables(), "==global_variables==")
if output_type == "sess":
return {
"train": {
"loss": loss,
"logits": logits,
"train_op": train_op
},
"hooks": training_hooks
}
elif output_type == "estimator":
return estimator_spec
elif mode == tf.estimator.ModeKeys.PREDICT:
print(logits.get_shape(), "===logits shape===")
label_weights = tf.cast(features['label_weights'], tf.int32)
label_seq_length = tf.reduce_sum(label_weights, axis=-1)
decode_tags, best_score = tf.contrib.crf.crf_decode(
logits, transition_params, label_seq_length)
_, hooks = model_io_fn.get_ema_hooks(
None, None, kargs.get('params_moving_average_decay', 0.99),
scope, mode)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
'decode_tags': decode_tags,
"best_score": best_score,
"transition_params": transition_params,
"logits": logits
},
export_outputs={
"output":
tf.estimator.export.PredictOutput({
'decode_tags': decode_tags,
"best_score": best_score,
"transition_params": transition_params,
"logits": logits
})
},
prediction_hooks=[hooks])
return estimator_spec
elif mode == tf.estimator.ModeKeys.EVAL:
_, hooks = model_io_fn.get_ema_hooks(
None, None, kargs.get('params_moving_average_decay', 0.99),
scope, mode)
eval_hooks = []
if output_type == "sess":
return {
"eval": {
"per_example_loss": per_example_loss,
"logits": logits,
"loss": tf.reduce_mean(per_example_loss),
"feature": model.get_pooled_output()
}
}
elif output_type == "estimator":
eval_metric_ops = eval_logtis(logits, features, num_labels,
transition_params)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=eval_hooks)
return estimator_spec
else:
raise NotImplementedError()
def model_fn(features, labels, mode):
model = bert_encoder(model_config, features, labels,
mode, target, reuse=model_reuse)
label_ids = features["label_ids"]
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=model_reuse):
(loss,
per_example_loss,
logits) = classifier.classifier(model_config,
model.get_pooled_output(),
num_labels,
label_ids,
dropout_prob)
model_io_fn = model_io.ModelIO(model_io_config)
tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
if load_pretrained:
model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope)
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer_fn = optimizer.Optimizer(opt_config)
model_io_fn.print_params(tvars, string=", trainable params")
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer_fn.get_train_op(loss, tvars,
opt_config.init_lr,
opt_config.num_train_steps)
model_io_fn.set_saver()
if kargs.get("task_index", 1) == 0:
model_io_fn.get_hooks(kargs.get("checkpoint_dir", None),
kargs.get("num_storage_steps", 1000))
training_hooks = model_io_fn.checkpoint_hook
else:
training_hooks = []
if len(optimizer_fn.distributed_hooks) >= 1:
training_hooks.extend(optimizer_fn.distributed_hooks)
print(training_hooks)
estimator_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=loss, train_op=train_op,
training_hooks=training_hooks)
if output_type == "sess":
return {
"train":{
"loss":loss,
"logits":logits,
"train_op":train_op
},
"hooks":training_hooks
}
elif output_type == "estimator":
return estimator_spec
elif mode == tf.estimator.ModeKeys.PREDICT:
print(logits.get_shape(), "===logits shape===")
pred_label = tf.argmax(logits, axis=-1, output_type=tf.int32)
prob = tf.nn.softmax(logits)
max_prob = tf.reduce_max(prob, axis=-1)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
'pred_label':pred_label,
"max_prob":max_prob
},
export_outputs={
"output":tf.estimator.export.PredictOutput(
{
'pred_label':pred_label,
"max_prob":max_prob
}
)
}
)
return estimator_spec
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss,
logits,
label_ids):
"""Computes the loss and accuracy of the model."""
sentence_log_probs = tf.reshape(
logits, [-1, logits.shape[-1]])
sentence_predictions = tf.argmax(
logits, axis=-1, output_type=tf.int32)
sentence_labels = tf.reshape(label_ids, [-1])
sentence_accuracy = tf.metrics.accuracy(
labels=label_ids, predictions=sentence_predictions)
sentence_mean_loss = tf.metrics.mean(
values=per_example_loss)
sentence_f = tf_metrics.f1(label_ids,
sentence_predictions,
num_labels,
label_lst, average="macro")
eval_metric_ops = {
"f1": sentence_f,
"acc":sentence_accuracy
}
return eval_metric_ops
eval_metric_ops = metric_fn(
per_example_loss,
logits,
label_ids)
estimator_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops)
if output_type == "sess":
return {
"eval":{
"per_example_loss":per_example_loss,
"logits":logits,
"loss":tf.reduce_mean(per_example_loss)
}
}
elif output_type == "estimator":
return estimator_spec
else:
raise NotImplementedError()
def model_fn(features, labels, mode):
features['input_mask'] = tf.cast(
tf.not_equal(features['input_ids'], kargs.get('[PAD]', 0)),
tf.int64)
# for key in ['input_mask', 'input_ids', 'segment_ids']:
# features[key] = features[key][:, :274]
model = bert_encoder(model_config,
features,
labels,
mode,
target,
reuse=tf.AUTO_REUSE)
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=model_reuse):
(loss, per_example_loss,
logits) = multi_position_classifier(model_config, features,
model.get_sequence_output(),
num_labels, dropout_prob)
model_io_fn = model_io.ModelIO(model_io_config)
tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
try:
params_size = model_io_fn.count_params(model_config.scope)
print("==total params==", params_size)
except:
print("==not count params==")
print(tvars)
if load_pretrained == "yes":
model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope)
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer_fn = optimizer.Optimizer(opt_config)
model_io_fn.print_params(tvars, string=", trainable params")
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
print("==update_ops==", update_ops)
with tf.control_dependencies(update_ops):
train_op = optimizer_fn.get_train_op(
loss, tvars, opt_config.init_lr,
opt_config.num_train_steps, **kargs)
# train_op, hooks = model_io_fn.get_ema_hooks(train_op,
# tvars,
# kargs.get('params_moving_average_decay', 0.99),
# scope, mode,
# first_stage_steps=opt_config.num_warmup_steps,
# two_stage=True)
model_io_fn.set_saver()
train_metric_dict = train_metric_fn(logits, features, num_labels)
for key in train_metric_dict:
tf.summary.scalar(key, train_metric_dict[key])
tf.summary.scalar('learning_rate', optimizer_fn.learning_rate)
if kargs.get("task_index", 1) == 0 and kargs.get(
"run_config", None):
training_hooks = []
elif kargs.get("task_index", 1) == 0:
model_io_fn.get_hooks(kargs.get("checkpoint_dir", None),
kargs.get("num_storage_steps", 1000))
training_hooks = model_io_fn.checkpoint_hook
else:
training_hooks = []
if len(optimizer_fn.distributed_hooks) >= 1:
training_hooks.extend(optimizer_fn.distributed_hooks)
print(training_hooks, "==training_hooks==", "==task_index==",
kargs.get("task_index", 1))
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=train_op,
training_hooks=training_hooks)
print(tf.global_variables(), "==global_variables==")
if output_type == "sess":
return {
"train": {
"loss": loss,
"logits": logits,
"train_op": train_op
},
"hooks": training_hooks
}
elif output_type == "estimator":
return estimator_spec
elif mode == tf.estimator.ModeKeys.PREDICT:
print(logits.get_shape(), "===logits shape===")
pred_label = tf.argmax(logits, axis=-1, output_type=tf.int32)
prob = tf.nn.softmax(logits)
max_prob = tf.reduce_max(prob, axis=-1)
# _, hooks = model_io_fn.get_ema_hooks(None,
# None,
# kargs.get('params_moving_average_decay', 0.99),
# scope, mode)
hooks = []
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
'pred_label': pred_label,
"max_prob": max_prob
},
export_outputs={
"output":
tf.estimator.export.PredictOutput({
'pred_label': pred_label,
"max_prob": max_prob
})
},
prediction_hooks=hooks)
return estimator_spec
elif mode == tf.estimator.ModeKeys.EVAL:
_, hooks = model_io_fn.get_ema_hooks(
None, None, kargs.get('params_moving_average_decay', 0.99),
scope, mode)
eval_hooks = []
if output_type == "sess":
return {
"eval": {
"per_example_loss": per_example_loss,
"logits": logits,
"loss": tf.reduce_mean(per_example_loss),
"feature": model.get_pooled_output()
}
}
elif output_type == "estimator":
eval_metric_ops = eval_logtis(logits, features, num_labels)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops,
evaluation_hooks=eval_hooks)
return estimator_spec
else:
raise NotImplementedError()
def model_fn(features, labels, mode):
model_lst = []
for index, name in enumerate(target):
if index > 0:
reuse = True
else:
reuse = model_reuse
model_lst.append(bert_encoder(model_config, features,
labels, mode, name, reuse=reuse))
label_ids = features["label_ids"]
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
else:
dropout_prob = 0.0
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
with tf.variable_scope(scope, reuse=reuse):
seq_output_lst = [model.get_pooled_output() for model in model_lst]
[loss,
per_example_loss,
logits] = classifier.order_classifier(
model_config, seq_output_lst,
num_labels, label_ids,
dropout_prob,ratio_weight)
tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
if load_pretrained:
model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope)
model_io_fn.set_saver(var_lst=tvars)
if mode == tf.estimator.ModeKeys.TRAIN:
model_io_fn.print_params(tvars, string=", trainable params")
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = optimizer_fn.get_train_op(loss, tvars,
opt_config.init_lr,
opt_config.num_train_steps)
estimator_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=loss, train_op=train_op)
return {
"estimator_spec":estimator_spec,
"train":{
"loss":loss,
"logits":logits,
"train_op":train_op
}
}
elif mode == tf.estimator.ModeKeys.PREDICT:
print(logits.get_shape(), "===logits shape===")
pred_label = tf.argmax(logits, axis=-1, output_type=tf.int32)
prob = tf.nn.softmax(logits)
max_prob = tf.reduce_max(prob, axis=-1)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
'pred_label':pred_label,
"max_prob":max_prob
},
export_outputs={
"output":tf.estimator.export.PredictOutput(
{
'pred_label':pred_label,
"max_prob":max_prob
}
)
}
)
return {
"estimator_spec":estimator_spec
}
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(per_example_loss,
logits,
label_ids):
"""Computes the loss and accuracy of the model."""
sentence_log_probs = tf.reshape(
logits, [-1, logits.shape[-1]])
sentence_predictions = tf.argmax(
logits, axis=-1, output_type=tf.int32)
sentence_labels = tf.reshape(label_ids, [-1])
sentence_accuracy = tf.metrics.accuracy(
labels=label_ids, predictions=sentence_predictions)
sentence_mean_loss = tf.metrics.mean(
values=per_example_loss)
sentence_f = tf_metrics.f1(label_ids,
sentence_predictions,
num_labels,
label_lst, average="macro")
eval_metric_ops = {
"f1": sentence_f,
"loss": sentence_mean_loss,
"acc":sentence_accuracy
}
return eval_metric_ops
eval_metric_ops = metric_fn(
per_example_loss,
logits,
label_ids)
estimator_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=loss,
eval_metric_ops=eval_metric_ops)
return {
"estimator_spec":estimator_spec,
"eval":{
"per_example_loss":per_example_loss,
"logits":logits
}
}
else:
raise NotImplementedError()