def model_fn(features, labels, mode):
task_type = kargs.get("task_type", "cls")
num_task = kargs.get('num_task', 1)
temp = kargs.get('temp', 0.1)
print("==task_type==", task_type)
model_io_fn = model_io.ModelIO(model_io_config)
label_ids = tf.cast(features["{}_label_ids".format(task_type)], dtype=tf.int32)
if mode == tf.estimator.ModeKeys.TRAIN:
dropout_prob = model_config.dropout_prob
is_training = True
else:
dropout_prob = 0.0
is_training = False
if model_io_config.fix_lm == True:
scope = model_config.scope + "_finetuning"
else:
scope = model_config.scope
if kargs.get("get_pooled_output", "pooled_output") == "pooled_output":
pooled_feature = model.get_pooled_output()
elif kargs.get("get_pooled_output", "task_output") == "task_output":
pooled_feature_dict = model.get_task_output()
pooled_feature = pooled_feature_dict['pooled_feature']
shape_list = bert_utils.get_shape_list(pooled_feature_dict['feature_a'],
expected_rank=[2])
batch_size = shape_list[0]
if kargs.get('apply_head_proj', False):
with tf.variable_scope(scope+"/head_proj", reuse=tf.AUTO_REUSE):
feature_a = simclr_utils.projection_head(pooled_feature_dict['feature_a'],
is_training,
head_proj_dim=128,
num_nlh_layers=1,
head_proj_mode='nonlinear',
name='head_contrastive')
pooled_feature_dict['feature_a'] = feature_a
with tf.variable_scope(scope+"/head_proj", reuse=tf.AUTO_REUSE):
feature_b = simclr_utils.projection_head(pooled_feature_dict['feature_b'],
is_training,
head_proj_dim=128,
num_nlh_layers=1,
head_proj_mode='nonlinear',
name='head_contrastive')
pooled_feature_dict['feature_b'] = feature_b
tf.logging.info("****** apply contrastive feature projection *******")
else:
feature_a = pooled_feature_dict['feature_a']
feature_b = pooled_feature_dict['feature_b']
tf.logging.info("****** not apply projection *******")
loss_mask = tf.cast(features["{}_loss_multipiler".format(task_type)], tf.float32)
if kargs.get('merge_mode', 'all') == 'all':
input_ids = tf.concat([features['input_ids_a'], features['input_ids_b']], axis=0)
hidden_repres = tf.concat([feature_a, feature_b], axis=0)
sent_repres = tf.concat([pooled_feature_dict['sent_repres_a'], pooled_feature_dict['sent_repres_b']], axis=0)
tf.logging.info("****** double batch *******")
else:
input_ids = features['input_ids_b']
hidden_repres = feature_b
sent_repres = pooled_feature_dict['sent_repres_b']
tf.logging.info("****** single batch b *******")
sequence_mask = tf.to_float(tf.not_equal(input_ids,
kargs.get('[PAD]', 0)))
with tf.variable_scope("vae/connect", reuse=tf.AUTO_REUSE):
with tf.variable_scope("z_mean"):
z_mean = tf.layers.dense(
hidden_repres,
128,
use_bias=None,
activation=None,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))
bn_z_mean = vae_utils.mean_normalize_scale(z_mean,
is_training,
"bn_mean",
tau=0.5,
reuse=tf.AUTO_REUSE,
**kargs)
with tf.variable_scope("z_std"):
z_std = tf.layers.dense(
hidden_repres,
128,
use_bias=True,
activation=tf.nn.relu,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))
bn_z_std = vae_utils.std_normalize_scale(z_std,
is_training,
"bn_std",
tau=0.5,
reuse=tf.AUTO_REUSE,
**kargs)
gaussian_noise = vae_utils.hidden_sampling(bn_z_mean, bn_z_std, **kargs)
sent_repres_shape = bert_utils.get_shape_list(sent_repres, expected_rank=[3])
with tf.variable_scope("vae/projection"):
gaussian_noise = tf.layers.dense(
gaussian_noise,
sent_repres_shape[-1],
use_bias=None,
activation=None,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01))
sent_repres += tf.expand_dims(gaussian_noise, axis=1)
with tf.variable_scope("vae/decoder", reuse=tf.AUTO_REUSE):
sequence_output = dgcnn_utils.dgcnn(
sent_repres,
sequence_mask,
num_layers=model_config['cnn_num_layers'],
dilation_rates=model_config.get('cnn_dilation_rates', [1,2]),
strides=model_config.get('cnn_dilation_rates', [1,1]),
num_filters=model_config.get('cnn_num_filters', [128,128]),
kernel_sizes=model_config.get('cnn_filter_sizes', [3,3]),
is_training=is_training,
scope_name="textcnn_encoder/textcnn/forward",
reuse=tf.AUTO_REUSE,
activation=tf.nn.relu,
is_casual=model_config['is_casual'],
padding=model_config.get('padding', 'same')
)
sequence_output_logits = model.build_other_output_logits(sequence_output, reuse=tf.AUTO_REUSE)
resc_loss = vae_utils.reconstruction_loss(sequence_output_logits,
input_ids,
name="decoder_resc",
use_tpu=False)
kl_loss = vae_utils.kl_loss(bn_z_mean, bn_z_std,
opt_config.get('num_train_steps', 10000),
name="kl_div",
use_tpu=False,
kl_anneal="kl_anneal")
loss = resc_loss + kl_loss
task_loss = loss
params_size = model_io_fn.count_params(model_config.scope)
print("==total encoder params==", params_size)
if kargs.get("feature_distillation", True):
universal_feature_a = features.get("input_ids_a_features", None)
universal_feature_b = features.get("input_ids_b_features", None)
if universal_feature_a is None or universal_feature_b is None:
tf.logging.info("****** not apply feature distillation *******")
feature_loss = tf.constant(0.0)
else:
feature_a = pooled_feature_dict['feature_a']
feature_a_shape = bert_utils.get_shape_list(feature_a, expected_rank=[2,3])
pretrain_feature_a_shape = bert_utils.get_shape_list(universal_feature_a, expected_rank=[2,3])
if feature_a_shape[-1] != pretrain_feature_a_shape[-1]:
with tf.variable_scope(scope+"/feature_proj", reuse=tf.AUTO_REUSE):
proj_feature_a = tf.layers.dense(feature_a, pretrain_feature_a_shape[-1])
# with tf.variable_scope(scope+"/feature_rec", reuse=tf.AUTO_REUSE):
# proj_feature_a_rec = tf.layers.dense(proj_feature_a, feature_a_shape[-1])
# loss += tf.reduce_mean(tf.reduce_sum(tf.square(proj_feature_a_rec-feature_a), axis=-1))/float(num_task)
tf.logging.info("****** apply auto-encoder for feature compression *******")
else:
proj_feature_a = feature_a
feature_a_norm = tf.stop_gradient(tf.sqrt(tf.reduce_sum(tf.pow(proj_feature_a, 2), axis=-1, keepdims=True))+1e-20)
proj_feature_a /= feature_a_norm
feature_b = pooled_feature_dict['feature_b']
if feature_a_shape[-1] != pretrain_feature_a_shape[-1]:
with tf.variable_scope(scope+"/feature_proj", reuse=tf.AUTO_REUSE):
proj_feature_b = tf.layers.dense(feature_b, pretrain_feature_a_shape[-1])
# with tf.variable_scope(scope+"/feature_rec", reuse=tf.AUTO_REUSE):
# proj_feature_b_rec = tf.layers.dense(proj_feature_b, feature_a_shape[-1])
# loss += tf.reduce_mean(tf.reduce_sum(tf.square(proj_feature_b_rec-feature_b), axis=-1))/float(num_task)
tf.logging.info("****** apply auto-encoder for feature compression *******")
else:
proj_feature_b = feature_b
feature_b_norm = tf.stop_gradient(tf.sqrt(tf.reduce_sum(tf.pow(proj_feature_b, 2), axis=-1, keepdims=True))+1e-20)
proj_feature_b /= feature_b_norm
feature_a_distillation = tf.reduce_mean(tf.square(universal_feature_a-proj_feature_a), axis=-1)
feature_b_distillation = tf.reduce_mean(tf.square(universal_feature_b-proj_feature_b), axis=-1)
feature_loss = tf.reduce_mean((feature_a_distillation + feature_b_distillation)/2.0)/float(num_task)
loss += feature_loss
tf.logging.info("****** apply prertained feature distillation *******")
if kargs.get("embedding_distillation", True):
word_embed = model.emb_mat
random_embed_shape = bert_utils.get_shape_list(word_embed, expected_rank=[2,3])
print("==random_embed_shape==", random_embed_shape)
pretrained_embed = kargs.get('pretrained_embed', None)
if pretrained_embed is None:
tf.logging.info("****** not apply prertained feature distillation *******")
embed_loss = tf.constant(0.0)
else:
pretrain_embed_shape = bert_utils.get_shape_list(pretrained_embed, expected_rank=[2,3])
print("==pretrain_embed_shape==", pretrain_embed_shape)
if random_embed_shape[-1] != pretrain_embed_shape[-1]:
with tf.variable_scope(scope+"/embedding_proj", reuse=tf.AUTO_REUSE):
proj_embed = tf.layers.dense(word_embed, pretrain_embed_shape[-1])
else:
proj_embed = word_embed
embed_loss = tf.reduce_mean(tf.reduce_mean(tf.square(proj_embed-pretrained_embed), axis=-1))/float(num_task)
loss += embed_loss
tf.logging.info("****** apply prertained feature distillation *******")
if mode == tf.estimator.ModeKeys.TRAIN:
multi_task_config = kargs.get("multi_task_config", {})
if multi_task_config.get(task_type, {}).get("lm_augumentation", False):
print("==apply lm_augumentation==")
masked_lm_positions = features["masked_lm_positions"]
masked_lm_ids = features["masked_lm_ids"]
masked_lm_weights = features["masked_lm_weights"]
(masked_lm_loss,
masked_lm_example_loss,
masked_lm_log_probs) = pretrain.get_masked_lm_output(
model_config,
model.get_sequence_output(),
model.get_embedding_table(),
masked_lm_positions,
masked_lm_ids,
masked_lm_weights,
reuse=model_reuse)
masked_lm_loss_mask = tf.expand_dims(loss_mask, -1) * tf.ones((1, multi_task_config[task_type]["max_predictions_per_seq"]))
masked_lm_loss_mask = tf.reshape(masked_lm_loss_mask, (-1, ))
masked_lm_label_weights = tf.reshape(masked_lm_weights, [-1])
masked_lm_loss_mask *= tf.cast(masked_lm_label_weights, tf.float32)
masked_lm_example_loss *= masked_lm_loss_mask# multiply task_mask
masked_lm_loss = tf.reduce_sum(masked_lm_example_loss) / (1e-10+tf.reduce_sum(masked_lm_loss_mask))
loss += multi_task_config[task_type]["masked_lm_loss_ratio"]*masked_lm_loss
masked_lm_label_ids = tf.reshape(masked_lm_ids, [-1])
print(masked_lm_log_probs.get_shape(), "===masked lm log probs===")
print(masked_lm_label_ids.get_shape(), "===masked lm ids===")
print(masked_lm_label_weights.get_shape(), "===masked lm mask===")
lm_acc = build_accuracy(masked_lm_log_probs, masked_lm_label_ids, masked_lm_loss_mask)
if kargs.get("task_invariant", "no") == "yes":
print("==apply task adversarial training==")
with tf.variable_scope(scope+"/dann_task_invariant", reuse=model_reuse):
(_,
task_example_loss,
task_logits) = distillation_utils.feature_distillation(model.get_pooled_output(),
1.0,
features["task_id"],
kargs.get("num_task", 7),
dropout_prob,
True)
masked_task_example_loss = loss_mask * task_example_loss
masked_task_loss = tf.reduce_sum(masked_task_example_loss) / (1e-10+tf.reduce_sum(loss_mask))
loss += kargs.get("task_adversarial", 1e-2) * masked_task_loss
tvars = model_io_fn.get_params(model_config.scope,
not_storage_params=not_storage_params)
vae_tvars = model_io_fn.get_params("vae",
not_storage_params=not_storage_params)
if mode == tf.estimator.ModeKeys.TRAIN:
multi_task_config = kargs.get("multi_task_config", {})
if multi_task_config.get(task_type, {}).get("lm_augumentation", False):
print("==apply lm_augumentation==")
masked_lm_pretrain_tvars = model_io_fn.get_params("cls/predictions",
not_storage_params=not_storage_params)
tvars.extend(masked_lm_pretrain_tvars)
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":
[assignment_map,
initialized_variable_names] = model_io_utils.get_assigment_map_from_checkpoint(
tvars,
init_checkpoint,
exclude_scope="")
[assignment_map_vae,
initialized_variable_names_vae] = model_io_utils.get_assigment_map_from_checkpoint(
vae_tvars,
init_checkpoint,
exclude_scope="vae/decoder")
assignment_map.update(assignment_map_vae)
initialized_variable_names.update(initialized_variable_names_vae)
model_io_utils.init_pretrained(assignment_map, initialized_variable_names,
tvars+vae_tvars, init_checkpoint)
if mode == tf.estimator.ModeKeys.TRAIN:
train_metric_dict = train_metric(input_ids,
sequence_output_logits,
**kargs)
return_dict = {
"loss":loss,
"tvars":tvars+vae_tvars
}
return_dict["perplexity"] = train_metric_dict['perplexity']
return_dict["token_acc"] = train_metric_dict['token_acc']
return_dict["kl_div"] = kl_loss
if kargs.get("task_invariant", "no") == "yes":
return_dict["{}_task_loss".format(task_type)] = masked_task_loss
task_acc = build_accuracy(task_logits, features["task_id"], loss_mask)
return_dict["{}_task_acc".format(task_type)] = task_acc
if multi_task_config.get(task_type, {}).get("lm_augumentation", False):
return_dict["{}_masked_lm_loss".format(task_type)] = masked_lm_loss
return_dict["{}_masked_lm_acc".format(task_type)] = lm_acc
if kargs.get("embedding_distillation", True):
return_dict["embed_loss"] = embed_loss*float(num_task)
else:
return_dict["embed_loss"] = task_loss
if kargs.get("feature_distillation", True):
return_dict["feature_loss"] = feature_loss*float(num_task)
else:
return_dict["feature_loss"] = task_loss
return_dict["task_loss"] = task_loss
return return_dict
elif mode == tf.estimator.ModeKeys.EVAL:
eval_dict = {
"loss":loss,
"logits":logits,
"feature":model.get_pooled_output()
}
if kargs.get("adversarial", "no") == "adversarial":
eval_dict["task_logits"] = task_logits
return eval_dict
def model_fn(features, labels, mode):
model_api = model_zoo(model_config)
model = model_api(model_config,
features,
labels,
mode,
target,
reuse=model_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
(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=model_reuse)
masked_lm_positions = features["masked_lm_positions"]
masked_lm_ids = features["masked_lm_ids"]
masked_lm_weights = features["masked_lm_weights"]
(masked_lm_loss, masked_lm_example_loss,
masked_lm_log_probs) = pretrain.get_masked_lm_output(
model_config,
model.get_sequence_output(),
model.get_embedding_table(),
masked_lm_positions,
masked_lm_ids,
masked_lm_weights,
reuse=model_reuse)
loss = model_config.lm_ratio * masked_lm_loss + model_config.nsp_ratio * nsp_loss
model_io_fn = model_io.ModelIO(model_io_config)
if mode == tf.estimator.ModeKeys.TRAIN:
pretrained_tvars = model_io_fn.get_params(
model_config.scope, not_storage_params=not_storage_params)
lm_pretrain_tvars = model_io_fn.get_params(
"cls", not_storage_params=not_storage_params)
pretrained_tvars.extend(lm_pretrain_tvars)
optimizer_fn = optimizer.Optimizer(opt_config)
if load_pretrained:
model_io_fn.load_pretrained(pretrained_tvars,
init_checkpoint,
exclude_scope=exclude_scope)
tvars = pretrained_tvars
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)
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(
masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids,
masked_lm_weights, nsp_per_example_loss, nsp_log_prob,
features['next_sentence_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)
if output_type == "sess":
return {
"train": {
"loss": loss,
"nsp_log_pro": nsp_log_prob,
"train_op": train_op,
"masked_lm_loss": masked_lm_loss,
"next_sentence_loss": nsp_loss,
"masked_lm_log_pro": masked_lm_log_probs
},
"hooks": training_hooks
}
elif output_type == "estimator":
return estimator_spec
elif mode == tf.estimator.ModeKeys.PREDICT:
def prediction_fn(logits):
predictions = {
"nsp_classes":
tf.argmax(input=nsp_log_prob, axis=1),
"nsp_probabilities":
tf.exp(nsp_log_prob, name="nsp_softmax"),
"masked_vocab_classes":
tf.argmax(input=masked_lm_log_probs, axis=1),
"masked_probabilities":
tf.exp(masked_lm_log_probs, name='masked_softmax')
}
return predictions
predictions = prediction_fn(logits)
estimator_spec = tf.estimator.EstimatorSpec(
mode=mode,
predictions=predictions,
export_outputs={
"output": tf.estimator.export.PredictOutput(predictions)
})
return estimator_spec
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(masked_lm_example_loss, masked_lm_log_probs,
masked_lm_ids, masked_lm_weights,
next_sentence_example_loss, next_sentence_log_probs,
next_sentence_labels):
"""Computes the loss and accuracy of the model."""
masked_lm_log_probs = tf.reshape(
masked_lm_log_probs, [-1, masked_lm_log_probs.shape[-1]])
masked_lm_predictions = tf.argmax(masked_lm_log_probs,
axis=-1,
output_type=tf.int32)
masked_lm_example_loss = tf.reshape(masked_lm_example_loss,
[-1])
masked_lm_ids = tf.reshape(masked_lm_ids, [-1])
masked_lm_weights = tf.reshape(masked_lm_weights, [-1])
masked_lm_accuracy = tf.metrics.accuracy(
labels=masked_lm_ids,
predictions=masked_lm_predictions,
weights=masked_lm_weights)
masked_lm_mean_loss = tf.metrics.mean(
values=masked_lm_example_loss, weights=masked_lm_weights)
next_sentence_log_probs = tf.reshape(
next_sentence_log_probs,
[-1, next_sentence_log_probs.shape[-1]])
next_sentence_predictions = tf.argmax(next_sentence_log_probs,
axis=-1,
output_type=tf.int32)
next_sentence_labels = tf.reshape(next_sentence_labels, [-1])
next_sentence_accuracy = tf.metrics.accuracy(
labels=next_sentence_labels,
predictions=next_sentence_predictions)
next_sentence_mean_loss = tf.metrics.mean(
values=next_sentence_example_loss)
return {
"masked_lm_accuracy": masked_lm_accuracy,
"masked_lm_loss": masked_lm_mean_loss,
"next_sentence_accuracy": next_sentence_accuracy,
"next_sentence_loss": next_sentence_mean_loss
}
if output_type == "sess":
return {
"eval": {
"nsp_log_prob": nsp_log_prob,
"masked_lm_log_prob": masked_lm_log_probs,
"nsp_loss": nsp_loss,
"masked_lm_loss": masked_lm_loss,
"feature": model.get_pooled_output()
}
}
elif output_type == "estimator":
eval_metric_ops = metric_fn(masked_lm_example_loss,
masked_lm_log_probs, masked_lm_ids,
masked_lm_weights,
nsp_per_example_loss, nsp_log_prob,
features['next_sentence_labels'])
_, hooks = model_io_fn.get_ema_hooks(
None, None, kargs.get('params_moving_average_decay', 0.99),
scope, mode)
eval_hooks = [hooks] if hooks else []
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()