def model_fn(features, labels, mode, params):
model_api = model_zoo(model_config)
model = model_api(model_config, features, labels,
mode, target, reuse=tf.AUTO_REUSE,
**kargs)
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
(_,
_,
nsp_log_prob) = pretrain.get_next_sentence_output(model_config,
model.get_pooled_output(),
features['next_sentence_labels'],
reuse=tf.AUTO_REUSE)
with tf.variable_scope('cls/seq_predictions', reuse=tf.AUTO_REUSE):
(_,
logits,
_) = classifier(model_config,
model.get_sequence_output(),
features['input_ori_ids'],
features['input_ids'],
features['input_mask'],
2,
dropout_prob)
# ,
# loss='focal_loss')
# loss += 0.0 * nsp_loss
model_io_fn = model_io.ModelIO(model_io_config)
pretrained_tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
lm_seq_prediction_tvars = model_io_fn.get_params("cls/seq_predictions",
not_storage_params=not_storage_params)
lm_pretrain_tvars = model_io_fn.get_params("cls/seq_relationship",
not_storage_params=not_storage_params)
pretrained_tvars.extend(lm_seq_prediction_tvars)
pretrained_tvars.extend(lm_pretrain_tvars)
tvars = pretrained_tvars
print('==discriminator parameters==', tvars)
if load_pretrained == "yes":
use_tpu = 1 if kargs.get('use_tpu', False) else 0
scaffold_fn = model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope,
use_tpu=use_tpu)
else:
scaffold_fn = None
if mode == tf.estimator.ModeKeys.PREDICT:
mask = tf.cast(tf.expand_dims(features['input_mask'], axis=-1), tf.float32)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions={
"probs":tf.nn.softmax(logits)*mask
},
export_outputs={
"output":tf.estimator.export.PredictOutput(
{
"probs":tf.nn.softmax(logits)*mask
}
)
}
)
return estimator_spec
def model_fn(features, labels, mode, params):
model_api = model_zoo(model_config)
model = model_api(model_config, features, labels,
mode, target, reuse=tf.AUTO_REUSE,
**kargs)
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
(nsp_loss,
nsp_per_example_loss,
nsp_log_prob) = pretrain.get_next_sentence_output(model_config,
model.get_pooled_output(),
features['next_sentence_labels'],
reuse=tf.AUTO_REUSE)
with tf.variable_scope('cls/seq_predictions', reuse=tf.AUTO_REUSE):
(loss,
logits,
per_example_loss) = classifier(model_config,
model.get_sequence_output(),
features['input_ori_ids'],
features['ori_input_ids'],
features['input_mask'],
2,
dropout_prob,
ori_sampled_ids=features.get('ori_sampled_ids', None),
use_tpu=kargs.get('use_tpu', True))
tf.add_to_collection("discriminator_loss", loss)
loss += 0.0 * nsp_loss
model_io_fn = model_io.ModelIO(model_io_config)
pretrained_tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
lm_seq_prediction_tvars = model_io_fn.get_params("cls/seq_predictions",
not_storage_params=not_storage_params)
lm_pretrain_tvars = model_io_fn.get_params("cls/seq_relationship",
not_storage_params=not_storage_params)
pretrained_tvars.extend(lm_seq_prediction_tvars)
pretrained_tvars.extend(lm_pretrain_tvars)
tvars = pretrained_tvars
print('==discriminator parameters==', tvars)
if load_pretrained == "yes":
use_tpu = 1 if kargs.get('use_tpu', False) else 0
scaffold_fn = model_io_fn.load_pretrained(tvars,
init_checkpoint,
exclude_scope=exclude_scope,
use_tpu=use_tpu,
restore_var_name=model_config.get('restore_var_name', []))
else:
scaffold_fn = None
return_dict = {
"loss":loss,
"logits":logits,
"tvars":tvars,
"model":model,
"per_example_loss":per_example_loss
}
return return_dict