def create_model(args, pyreader_name, ernie_config):
src_ids = fluid.layers.data(name='1',
shape=[-1, args.max_seq_len, 1],
dtype='int64')
sent_ids = fluid.layers.data(name='2',
shape=[-1, args.max_seq_len, 1],
dtype='int64')
pos_ids = fluid.layers.data(name='3',
shape=[-1, args.max_seq_len, 1],
dtype='int64')
task_ids = fluid.layers.data(name='4',
shape=[-1, args.max_seq_len, 1],
dtype='int64')
input_mask = fluid.layers.data(name='5',
shape=[-1, args.max_seq_len, 1],
dtype='float32')
seq_lens = fluid.layers.data(name='8', shape=[-1], dtype='int64')
pyreader = fluid.io.DataLoader.from_generator(
feed_list=[src_ids, sent_ids, pos_ids, task_ids, input_mask, seq_lens],
capacity=70,
iterable=False)
ernie = ErnieModel(src_ids=src_ids,
position_ids=pos_ids,
sentence_ids=sent_ids,
input_mask=input_mask,
config=ernie_config)
enc_out = ernie.get_sequence_output()
unpad_enc_out = fluid.layers.sequence_unpad(enc_out, length=seq_lens)
cls_feats = ernie.get_pooled_output()
# set persistable = True to avoid memory opimizing
enc_out.persistable = True
unpad_enc_out.persistable = True
cls_feats.persistable = True
graph_vars = {
"cls_embeddings": cls_feats,
"top_layer_embeddings": unpad_enc_out,
}
return pyreader, graph_vars
def create_model(args, pyreader_name, ernie_config, is_prediction=False):
pyreader = fluid.layers.py_reader(capacity=50,
shapes=[[-1, args.max_seq_len, 1],
[-1, args.max_seq_len, 1],
[-1, args.max_seq_len, 1],
[-1, args.max_seq_len, 1],
[-1, 1], [-1, 1], [3, 1], [3]],
dtypes=[
'int64', 'int64', 'int64', 'float32',
'int64', 'int64', 'int64', 'int64'
],
lod_levels=[0, 0, 0, 0, 0, 0, 0, 0],
name=pyreader_name,
use_double_buffer=True)
(src_ids, sent_ids, pos_ids, input_mask, labels, qids, mlm_mask_label,
mlm_mask_pos) = fluid.layers.read_file(pyreader)
ernie = ErnieModel(src_ids=src_ids,
position_ids=pos_ids,
sentence_ids=sent_ids,
input_mask=input_mask,
config=ernie_config,
use_fp16=args.use_fp16)
cls_feats = ernie.get_pooled_output()
cls_feats = fluid.layers.dropout(x=cls_feats,
dropout_prob=0.1,
dropout_implementation="upscale_in_train")
logits = fluid.layers.fc(
input=cls_feats,
size=args.num_labels,
param_attr=fluid.ParamAttr(
name="cls_out_w",
initializer=fluid.initializer.TruncatedNormal(scale=0.02)),
bias_attr=fluid.ParamAttr(name="cls_out_b",
initializer=fluid.initializer.Constant(0.)))
ernie.get_pretraining_output(mask_label=mlm_mask_label,
mask_pos=mlm_mask_pos,
labels=labels)
if is_prediction:
probs = fluid.layers.softmax(logits)
feed_targets_name = [
src_ids.name, pos_ids.name, sent_ids.name, input_mask.name
]
return pyreader, probs, feed_targets_name
ce_loss, probs = fluid.layers.softmax_with_cross_entropy(
logits=logits, label=labels, return_softmax=True)
loss = fluid.layers.mean(x=ce_loss)
if args.use_fp16 and args.loss_scaling > 1.0:
loss *= args.loss_scaling
num_seqs = fluid.layers.create_tensor(dtype='int64')
accuracy = fluid.layers.accuracy(input=probs, label=labels, total=num_seqs)
graph_vars = {
"loss": loss,
"probs": probs,
"accuracy": accuracy,
"labels": labels,
"num_seqs": num_seqs,
"qids": qids
}
for k, v in graph_vars.items():
v.persistable = True
return pyreader, graph_vars