def model_fn(features, mode, params, run_config):
ernie = ErnieModelForSequenceClassification(params, name='')
if mode is not propeller.RunMode.TRAIN:
ernie.eval()
else:
ernie.train()
metrics, loss = None, None
if mode is propeller.RunMode.PREDICT:
src_ids, sent_ids = features
_, logits = ernie(src_ids, sent_ids)
predictions = [
logits,
]
else:
src_ids, sent_ids, labels = features
if mode is propeller.RunMode.EVAL:
loss, logits = ernie(src_ids, sent_ids, labels=labels)
pred = logits.argmax(axis=1)
acc = propeller.metrics.Acc(labels, pred)
metrics = {'acc': acc}
predictions = [pred]
train_hooks = None
else:
loss, logits = ernie(src_ids, sent_ids, labels=labels)
lr_step_hook, loss_scale_coef = optimization(
loss=loss,
warmup_steps=int(run_config.max_steps *
params['warmup_proportion']),
num_train_steps=run_config.max_steps,
learning_rate=params['learning_rate'],
train_program=P.static.default_main_program(),
startup_prog=P.static.default_startup_program(),
use_fp16=args.use_amp,
weight_decay=params['weight_decay'],
scheduler="linear_warmup_decay",
)
scheduled_lr = P.static.default_main_program().global_block().var(
'learning_rate_0')
propeller.summary.scalar('lr', scheduled_lr)
predictions = [
logits,
]
train_hooks = [lr_step_hook]
return propeller.ModelSpec(loss=loss,
mode=mode,
metrics=metrics,
predictions=predictions,
train_hooks=train_hooks)
model.clear_gradients()
if global_step % args.save_steps == 0:
F.save_dygraph(model.state_dict(),
args.save_dir + '_%s' % global_step)
if global_step % args.eval_steps == 0 and step > 0:
y_true, y_pred = [], []
with FD.base._switch_tracer_mode_guard_(is_train=False):
model.eval()
for step, d in enumerate(
tqdm(dev_batch_data,
desc='evaluating %d' % epoch)):
ids, sids, labels = d
ids, sids, labels = FD.to_variable(
ids), FD.to_variable(sids), FD.to_variable(
labels)
loss, logits = model(ids, sids, labels=labels)
#print('\n'.join(map(str, logits.numpy().tolist())))
y_pred += L.argmax(logits, -1).numpy().tolist()
y_true += labels.numpy().tolist()
model.train()
if args.debug:
print(y_pred[:10], y_true[:10])
f1 = f1_score(y_true, y_pred)
print('f1 %.5f' % f1)
print(classification_report(y_true, y_pred))
if f1 > bst_f1:
F.save_dygraph(model.state_dict(), args.save_dir)
bst_f1 = f1
print('saving model with best f1: %.3f' % bst_f1)