def create_model(pyreader_name, ernie_config, task_group): """create_model""" ## get input shapes = [[bsz, args.max_seq_len, 1], [bsz, args.max_seq_len, 1], [bsz, args.max_seq_len, 1], [bsz, args.max_seq_len, 1], [bsz, args.max_seq_len, 1], [bsz, 1], [bsz, 1], [1], [bsz, 1], [bsz, 1], [bsz, 1]] names = [ "src_ids", "pos_ids", "sent_ids", "task_ids", "input_mask", "mask_label", "mask_pos", "lm_weight", "batch_mask", "loss_mask", "gather_idx" ] dtypes = [ "int64", "int64", "int64", "int64", "float32", "int64", "int64", "float32", "float32", "float32", "int64" ] cnt_general_input = len(shapes) for index, task in enumerate(task_group): shapes.extend([[bsz, 1], [1]]) names.extend(['task_label_' + str(index), 'task_weight_' + str(index)]) dtypes.extend(["int64", "float32"]) assert len(shapes) == len(names) == len( dtypes), "The three fields must have same size" inputs = [] for i in range(len(shapes)): inputs.append( fluid.layers.data(name=names[i], shape=shapes[i], dtype=dtypes[i], append_batch_size=False)) general_data, task_params = inputs[:cnt_general_input], inputs[ cnt_general_input:] src_ids, pos_ids, sent_ids, task_ids, input_mask, \ mask_label, mask_pos, lm_weight, batch_mask, loss_mask, gather_idx = general_data ## build graph ernie = ErnieModel(src_ids=src_ids, position_ids=pos_ids, sentence_ids=sent_ids, task_ids=task_ids, input_mask=input_mask, config=ernie_config, weight_sharing=args.weight_sharing, use_fp16=args.use_amp) mask_lm_loss = ernie.get_lm_output(mask_label, mask_pos) checkpoints = ernie.get_checkpoints() total_loss = mask_lm_loss * lm_weight graph_vars = [mask_lm_loss, lm_weight] index = 0 total_constract_loss = 0 for task in task_group: task_labels = task_params[index] task_weight = task_params[index + 1] task_loss, task_acc = ernie.get_task_output(task, task_labels, gather_idx) total_loss += task_loss * task_weight * task["loss_weight"] if task["constart"]: contract_loss = ernie.get_contrastive_loss(batch_mask, loss_mask) total_loss += contract_loss * task_weight total_constract_loss += contract_loss * task_weight graph_vars.extend([task_acc, task_weight]) index += 2 ## build output graph_vars.append(total_constract_loss) graph_vars.append(total_loss) #for var in graph_vars: # var.persistable = True fetch_vars = {"graph_vars": graph_vars, "checkpoints": checkpoints} return fetch_vars, names
def create_model(pyreader_name, ernie_config, task_group): """create_model""" src_ids = fluid.layers.data(name='src_ids', shape=[-1, args.max_seq_len, 1], dtype='int64') pos_ids = fluid.layers.data(name='pos_ids', shape=[-1, args.max_seq_len, 1], dtype='int64') sent_ids= fluid.layers.data(name='sent_ids', shape=[-1, args.max_seq_len, 1], dtype='int64') task_ids= fluid.layers.data(name='task_ids', shape=[-1, args.max_seq_len, 1], dtype='int64') input_mask = fluid.layers.data(name='input_mask', shape=[-1, args.max_seq_len, args.max_seq_len], dtype='float32') mask_label = fluid.layers.data(name='mask_label', shape=[-1, 1], dtype='int64') mask_pos = fluid.layers.data(name='mask_pos', shape=[-1, 1], dtype='int64') lm_weight = fluid.layers.data(name='lm_weight', shape=[1], dtype='float32', append_batch_size=False) batch_mask = fluid.layers.data(name='batch_mask', shape=[-1, 1], dtype='float32') loss_mask = fluid.layers.data(name="loss_mask", shape=[-1, 1], dtype='float32') gather_idx = fluid.layers.data(name="gather_idx", shape=[-1, 1], dtype='int64') task_params_all = [] for index, task in enumerate(task_group): name_label = 'task_label_' + str(index) name_weight = 'task_weight_' + str(index) task_label = fluid.layers.data(name=name_label, shape=[-1, 1], dtype='int64') task_weight = fluid.layers.data(name=name_weight, shape=[1], dtype='float32', append_batch_size=False) task_params_all.extend([task_label, task_weight]) fluid.reader.keep_data_loader_order(False) feed_list = [src_ids, pos_ids, sent_ids, task_ids, input_mask, \ mask_label, mask_pos, lm_weight, batch_mask, loss_mask, gather_idx] + task_params_all pyreader = fluid.io.DataLoader.from_generator( feed_list=feed_list, capacity=70, iterable=False) ernie = ErnieModel( src_ids=src_ids, position_ids=pos_ids, sentence_ids=sent_ids, task_ids=task_ids, input_mask=input_mask, config=ernie_config, weight_sharing=args.weight_sharing, use_fp16=args.use_amp) mask_lm_loss = ernie.get_lm_output(mask_label, mask_pos) checkpoints = ernie.get_checkpoints() total_loss = mask_lm_loss * lm_weight graph_vars = [mask_lm_loss, lm_weight] index = 11 total_constract_loss = 0 for task in task_group: task_labels = feed_list[index] task_weight = feed_list[index + 1] task_loss, task_acc = ernie.get_task_output(task, task_labels, gather_idx) total_loss += task_loss * task_weight * task["loss_weight"] if task["constart"]: contract_loss = ernie.get_contrastive_loss(batch_mask, loss_mask) total_loss += contract_loss * task_weight total_constract_loss += contract_loss * task_weight graph_vars.extend([task_acc, task_weight]) index += 2 graph_vars.append(total_constract_loss) graph_vars.append(total_loss) for var in graph_vars: var.persistable = True fetch_vars = {"graph_vars": graph_vars, "checkpoints": checkpoints} return pyreader, fetch_vars
def create_model(args, phase, micro_bsz, dp_sharding_rank, dp_worldsize, topo): if args.use_sop: from reader.pretraining_ds_ernie_full_sent import make_pretrain_dataset else: from reader.pretraining_ds_mlm import make_pretrain_dataset # mask_label, mask_pos for mlm, labels for sop if args.use_sop: input_fields = { 'names': ['src_ids', 'sent_ids', 'mask_label', 'mask_pos', 'labels'], 'shapes': [[-1, args.max_seq_len, 1], [-1, args.max_seq_len, 1], [-1, 1], [-1, 1], [-1, 1]], 'dtypes': ['int64', 'int64', 'int64', 'int64', 'int64'], 'lod_levels': [0, 0, 0, 0, 0], } else: input_fields = { 'names': ['src_ids', 'sent_ids', 'mask_label', 'mask_pos'], 'shapes': [[-1, args.max_seq_len, 1], [-1, args.max_seq_len, 1], [-1, 1], [-1, 1]], 'dtypes': ['int64', 'int64', 'int64', 'int64'], 'lod_levels': [0, 0, 0, 0], } with fluid.device_guard("gpu:0"): inputs = [ fluid.data(name=input_fields['names'][i], shape=input_fields['shapes'][i], dtype=input_fields['dtypes'][i], lod_level=input_fields['lod_levels'][i]) for i in range(len(input_fields['names'])) ] if args.use_sop: (src_ids, sent_ids, mask_label, mask_pos, labels) = inputs else: (src_ids, sent_ids, mask_label, mask_pos) = inputs train_file_list = glob.glob(args.data_dir + "/*") vocab = {} with open(args.vocab_file) as r: for line in r: lines = line.strip().split('\t') vocab[lines[0]] = int(lines[1]) log.debug("========= worker: {} of {} ==========".format( dp_sharding_rank, dp_worldsize)) data_reader = make_pretrain_dataset('pt', train_file_list, True, vocab, micro_bsz, len(vocab), args.max_seq_len, dp_sharding_rank, dp_worldsize) with fluid.device_guard("gpu:0"): data_loader = fluid.io.DataLoader.from_generator(feed_list=inputs, capacity=70, iterable=False) places = fluid.CUDAPlace(int(os.environ.get('FLAGS_selected_gpus', 0))) def data_gen(): yield from data_reader data_loader.set_batch_generator(data_gen, places) ernie_config = ErnieConfig(args.ernie_config_file)._config_dict ernie_config["preln"] = args.preln weight_sharing = (topo.mp.size == 1 and topo.pp.size == 1 ) # pp mp should not do weight sharing with fluid.device_guard("gpu:0"): ernie = ErnieModel(src_ids, sent_ids, ernie_config, weight_sharing=weight_sharing, topo=topo) checkpoints = ernie._checkpoints checkpoints.pop(-1) with fluid.device_guard(f'gpu:{args.num_pp-1}'): mask_lm_loss, mean_mask_lm_loss = ernie.get_lm_output( mask_label, mask_pos) total_loss = mean_mask_lm_loss if args.use_sop: sop_acc, mean_sop_loss = ernie.get_next_sentence_output(labels) total_loss += mean_sop_loss if topo.pp.size > 1: mask_lm_loss.persistable = True mean_mask_lm_loss.persistable = True # checkpoints.extend([mask_lm_loss.name, mean_mask_lm_loss.name]) if args.use_sop: mean_sop_loss.persistable = True sop_acc.persistable = True # checkpoints.extend([mean_sop_loss.name, sop_acc.name]) total_loss.persistable = True # checkpoints.append(total_loss.name) if args.use_sop: graph_vars = { 'data_loader': data_loader, 'mask_lm_loss': mask_lm_loss, 'mean_mask_lm_loss': mean_mask_lm_loss, 'sop_loss': mean_sop_loss, 'sop_acc': sop_acc, 'total_loss': total_loss, 'checkpoints': checkpoints } else: graph_vars = { 'data_loader': data_loader, 'mask_lm_loss': mask_lm_loss, 'mean_mask_lm_loss': mean_mask_lm_loss, 'total_loss': total_loss, 'checkpoints': checkpoints, } return graph_vars