def main(): #np.random.seed(0) X, y = sklearn.datasets.make_moons(500, noise=.5) skew = np.array([[.5, .5], [-.5, .5]]) y *= 2 y -= 1 #X = X.dot(skew) plt.scatter(X[:, 0], X[:, 1], s=40, c=y, cmap=plt.cm.Spectral) data = i_d.IndexedDataset(X, y) for epoch in range(50): print(f"epoch {epoch}") data.new_learner() def predict(data, x): return data.classify_matrix(x) plot_decision_boundary(lambda x: predict(data, x), X, y) plt.show()
def main(): parser = argparse.ArgumentParser() ## Required parameters parser.add_argument("--data_dir", default=None, type=str, required=True, help="The input data dir. Should contain the .tsv files (or other data files) for the task.") parser.add_argument("--bert_model", default=None, type=str, required=True, help="Bert pre-trained model selected in the list: bert-base-uncased, " "bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.") parser.add_argument("--task_name", default=None, type=str, required=True, help="The name of the task to train.") parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.") ## Other parameters parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after WordPiece tokenization. \n" "Sequences longer than this will be truncated, and sequences shorter \n" "than this will be padded.") parser.add_argument("--do_train", default=False, action='store_true', help="Whether to run training.") parser.add_argument("--do_eval", default=False, action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--do_lower_case", default=False, action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=16, type=int, help="Total batch size for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Total batch size for eval.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--num_train_epochs", default=1.0, type=float, help="Total number of training epochs to perform.") parser.add_argument("--warmup_proportion", default=0.1, type=float, help="Proportion of training to perform linear learning rate warmup for. " "E.g., 0.1 = 10%% of training.") parser.add_argument("--no_cuda", default=False, action='store_true', help="Whether not to use CUDA when available") parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument('--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument('--fp16', default=False, action='store_true', help="Whether to use 16-bit float precision instead of 32-bit") parser.add_argument('--loss_scale', type=float, default=0, help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n" "0 (default value): dynamic loss scaling.\n" "Positive power of 2: static loss scaling value.\n") parser.add_argument("--vec_file", default=None, type=str, required=True, help="File with embeddings") parser.add_argument("--use_lim_ents", default=None, type=str, required=True, help="Whether to use limited entities") args = parser.parse_args() master_ip = os.environ['MASTER_ADDR'] master_port = os.environ['MASTER_PORT'] world_size = os.environ['WORLD_SIZE'] rank = os.environ['RANK'] logger.info("Master node's IP Address: {}, port: {}, world_size: {}, rank: {}".format(master_ip, master_port, world_size, rank)) logger.info ("Local rank received by launch utility: {}".format(args.local_rank)) logger.info("Process is being blocked until all nodes are ready.") if args.local_rank == -1 or args.no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logger.info("All nodes ready, unblocking process.\n\n") logger.info("Global rank: {}, device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format( torch.distributed.get_rank(), device, n_gpu, bool(args.local_rank != -1), args.fp16)) if args.gradient_accumulation_steps < 1: raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format( args.gradient_accumulation_steps)) args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) if not args.do_train and not args.do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir)) os.makedirs(args.output_dir, exist_ok=True) task_name = args.task_name.lower() #check for limited ents lim_ents = [] lim_check = (args.use_lim_ents == "y") if lim_check: lim_ents = lim_ent_map(0,"kg_embeddings/dbp_eid_2_wd_eid.txt") logger.info("Limited entities flag is on. Count of unique entities considered: "+str(len(lim_ents))) vecs = [] vecs.append([0]*100) # CLS lineindex = 1 uid_map = {} logger.info("Reading embeddings file.") with open(args.vec_file, 'r') as fin: for line in fin: vec = line.strip().split('\t') #first element is unique id uniqid = int(vec[0]) #map line index to unique id uid_map[uniqid] = lineindex #increment line index lineindex = lineindex + 1 if (lim_check and (uniqid in lim_ents)) or not lim_check: vec = [float(x) for x in vec[1:101]] else: vec = vecs[0] vecs.append(vec) embed = torch.FloatTensor(vecs) embed = torch.nn.Embedding.from_pretrained(embed) #embed = torch.nn.Embedding(5041175, 100) logger.info("Shape of entity embedding: "+str(embed.weight.size())) del vecs train_data = None num_train_steps = None if args.do_train: # TODO import indexed_dataset from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler,BatchSampler import iterators #train_data = indexed_dataset.IndexedCachedDataset(args.data_dir) train_data = indexed_dataset.IndexedDataset(args.data_dir, fix_lua_indexing=True) if args.local_rank == -1: train_sampler = RandomSampler(train_data) else: train_sampler = DistributedSampler(train_data) train_sampler = BatchSampler(train_sampler, args.train_batch_size, True) def collate_fn(x): #logger.info("Data for collate\n" + str(x)) x = torch.LongTensor([xx for xx in x]) entity_idx = x[:, 4*args.max_seq_length:5*args.max_seq_length] #logger.info("Entity ids:\n" + str(x)) #fetch the line index for the unique id entarr = [] global keys_found global keys_missed for elarr in entity_idx: temp_arr = [] for uniqid in elarr: lval = uniqid.item() if lval in uid_map: temp_arr.append(uid_map[lval]) keys_found = keys_found + 1 else: temp_arr.append(0) keys_missed = keys_missed + 1 entarr.append(temp_arr) entarr = torch.LongTensor(entarr) #logger.info("Entity array for current line: "+str(entarr.numpy())) # Build candidate uniq_idx = np.unique(entarr.numpy()) ent_candidate = embed(torch.LongTensor(uniq_idx)) ent_candidate = ent_candidate.repeat([n_gpu, 1]) # build entity labels d = {} dd = [] for i, idx in enumerate(uniq_idx): d[idx] = i dd.append(idx) ent_size = len(uniq_idx)-1 def map(x): if x == -1 or x == 0: return 0 else: rnd = random.uniform(0, 1) if rnd < 0.05: return dd[random.randint(1, ent_size)] elif rnd < 0.2: return 0 else: return x ent_labels = entarr.clone() d[-1] = -1 ent_labels = ent_labels.apply_(lambda x: d[x]) entarr.apply_(map) ent_emb = embed(entarr) mask = entarr.clone() mask.apply_(lambda x: 0 if (x == -1 or x ==0) else 1) mask[:,0] = 1 return x[:,:args.max_seq_length], x[:,args.max_seq_length:2*args.max_seq_length], x[:,2*args.max_seq_length:3*args.max_seq_length], x[:,3*args.max_seq_length:4*args.max_seq_length], ent_emb, mask, x[:,6*args.max_seq_length:], ent_candidate, ent_labels train_iterator = iterators.EpochBatchIterator(train_data, collate_fn, train_sampler) num_train_steps = int( len(train_data) / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs) # Prepare model model, missing_keys = BertForPreTraining.from_pretrained(args.bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank)) model.to(device) if args.local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") model = DDP(model) elif n_gpu > 1: model = torch.nn.DataParallel(model) # Prepare optimizer param_optimizer = list(model.named_parameters()) no_linear = ['layer.2.output.dense_ent', 'layer.2.intermediate.dense_1', 'bert.encoder.layer.2.intermediate.dense_1_ent', 'layer.2.output.LayerNorm_ent'] no_linear = [x.replace('2', '11') for x in no_linear] param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in no_linear)] #param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in missing_keys)] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] t_total = num_train_steps if args.local_rank != -1: t_total = t_total // torch.distributed.get_world_size() if args.fp16: try: from apex.contrib.optimizers import FP16_Optimizer from apex.optimizers import FusedAdam from apex import amp except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") optimizer = FusedAdam(optimizer_grouped_parameters, lr=args.learning_rate, bias_correction=False) if args.loss_scale == 0: model, optimizer = amp.initialize(model, optimizer, opt_level="O2") # optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True) else: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", loss_scale=args.loss_scale) #logger.info(dir(optimizer)) #op_path = os.path.join(args.bert_model, "pytorch_op.bin") #optimizer.load_state_dict(torch.load(op_path)) else: optimizer = BertAdam(optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=t_total) global_step = 0 if args.do_train: logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_data)) logger.info(" Batch size = %d", args.train_batch_size) logger.info(" Num steps = %d", num_train_steps) model.train() import datetime fout = open(os.path.join(args.output_dir, "loss.{}".format(datetime.datetime.now())), 'w') for _ in trange(int(args.num_train_epochs), desc="Epoch"): tr_loss = 0 nb_tr_examples, nb_tr_steps = 0, 0 for step, batch in enumerate(tqdm(train_iterator.next_epoch_itr(), desc="Iteration")): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels = batch loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels) if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. original_loss = original_loss.mean() if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: try: from apex import amp except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() fout.write("{} {}\n".format(loss.item()*args.gradient_accumulation_steps, original_loss.item())) tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 if (step + 1) % args.gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses lr_this_step = args.learning_rate * warmup_linear(global_step/t_total, args.warmup_proportion) for param_group in optimizer.param_groups: param_group['lr'] = lr_this_step optimizer.step() optimizer.zero_grad() global_step += 1 #if global_step % 1000 == 0: # model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self # output_model_file = os.path.join(args.output_dir, "pytorch_model.bin_{}".format(global_step)) # torch.save(model_to_save.state_dict(), output_model_file) fout.close() logger.info("Saving data") # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model_"+str(torch.distributed.get_rank())+str(args.local_rank)+".bin") torch.save(model_to_save.state_dict(), output_model_file) logger.info("Training complete.\n Total number of entity matches in embeddings: ", keys_found, "\n Missed matches: ", keys_missed)
import numpy as np import indexed_dataset import torch from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler, BatchSampler import iterators import tensorflow as tf max_seq = 256 dataset = indexed_dataset.IndexedDataset('output/data_00') print(len(dataset)) train_sampler = BatchSampler(RandomSampler(dataset), 1, True) vocab_words = [] with tf.gfile.GFile("vocab.txt", 'r') as fin: for line in fin: vocab_words.append(line.strip()) d = {} with tf.gfile.GFile("entity2id.txt", 'r') as fin: fin.readline() for line in fin: ent, idx = line.split() d[int(idx)] = ent def collate_fn(x): x = torch.LongTensor([xx.numpy() for xx in x]) - 1 entity_idx = x[:, 4 * max_seq:5 * max_seq] # ent_emb = embed(entity_idx+1) # -1 -> 0 mask = x[:, 5 * max_seq:6 * max_seq] idx = np.unique(entity_idx.numpy())
def __init__(self, x, y): self.data = i_d.IndexedDataset(x, y) self.x = x self.y = y self.learners = [] self.alpha = []
def main(): parser = argparse.ArgumentParser() ## Required parameters parser.add_argument("--data_dir", default=None, type=str, required=True, help="The input data dir. Should contain the .tsv files (or other data files) for the task.") parser.add_argument("--bert_model", default=None, type=str, required=True, help="Bert pre-trained model selected in the list: bert-base-uncased, " "bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.") parser.add_argument("--task_name", default=None, type=str, required=True, help="The name of the task to train.") parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.") ## Other parameters parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after WordPiece tokenization. \n" "Sequences longer than this will be truncated, and sequences shorter \n" "than this will be padded.") parser.add_argument("--do_train", default=False, action='store_true', help="Whether to run training.") parser.add_argument("--do_eval", default=False, action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--do_lower_case", default=False, action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=32, type=int, help="Total batch size for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Total batch size for eval.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.") parser.add_argument("--warmup_proportion", default=0.1, type=float, help="Proportion of training to perform linear learning rate warmup for. " "E.g., 0.1 = 10%% of training.") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument('--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument('--loss_scale', type=float, default=0, help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n" "0 (default value): dynamic loss scaling.\n" "Positive power of 2: static loss scaling value.\n") args = parser.parse_args() args.local_rank = -1 device = torch.device("cpu") n_gpu = 0 logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format( device, n_gpu, bool(args.local_rank != -1), 'false')) if args.gradient_accumulation_steps < 1: raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format( args.gradient_accumulation_steps)) args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if not args.do_train and not args.do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir)) os.makedirs(args.output_dir, exist_ok=True) task_name = args.task_name.lower() vecs = [] vecs.append([0]*100) # CLS with open("kg_embed/entity2vec.vec", 'r') as fin: for line in fin: vec = line.strip().split('\t') vec = [float(x) for x in vec] vecs.append(vec) embed = torch.FloatTensor(vecs) embed = torch.nn.Embedding.from_pretrained(embed) #embed = torch.nn.Embedding(5041175, 100) logger.info("Shape of entity embedding: "+str(embed.weight.size())) del vecs train_data = None num_train_steps = None if args.do_train: import indexed_dataset from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler,BatchSampler import iterators #train_data = indexed_dataset.IndexedCachedDataset(args.data_dir) train_data = indexed_dataset.IndexedDataset(args.data_dir, fix_lua_indexing=True) train_sampler = RandomSampler(train_data) train_sampler = BatchSampler(train_sampler, args.train_batch_size, True) def collate_fn(x): x = torch.LongTensor([xx for xx in x]) entity_idx = x[:, 4*args.max_seq_length:5*args.max_seq_length] # Build candidate uniq_idx = np.unique(entity_idx.numpy()) ent_candidate = embed(torch.LongTensor(uniq_idx+1)) ent_candidate = ent_candidate.repeat([n_gpu, 1]) # build entity labels d = {} dd = [] for i, idx in enumerate(uniq_idx): d[idx] = i dd.append(idx) ent_size = len(uniq_idx)-1 def map(x): if x == -1: return -1 else: rnd = random.uniform(0, 1) if rnd < 0.05: return dd[random.randint(1, ent_size)] elif rnd < 0.2: return -1 else: return x ent_labels = entity_idx.clone() d[-1] = -1 ent_labels = ent_labels.apply_(lambda x: d[x]) entity_idx.apply_(map) ent_emb = embed(entity_idx+1) mask = entity_idx.clone() mask.apply_(lambda x: 0 if x == -1 else 1) mask[:,0] = 1 return x[:,:args.max_seq_length], x[:,args.max_seq_length:2*args.max_seq_length], x[:,2*args.max_seq_length:3*args.max_seq_length], x[:,3*args.max_seq_length:4*args.max_seq_length], ent_emb, mask, x[:,6*args.max_seq_length:], ent_candidate, ent_labels train_iterator = iterators.EpochBatchIterator(train_data, collate_fn, train_sampler) num_train_steps = int( len(train_data) / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs) # Prepare model model, missing_keys = BertForPreTraining.from_pretrained(args.bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(-1)) model.to(device) # Prepare optimizer param_optimizer = list(model.named_parameters()) no_linear = ['layer.2.output.dense_ent', 'layer.2.intermediate.dense_1', 'bert.encoder.layer.2.intermediate.dense_1_ent', 'layer.2.output.LayerNorm_ent'] no_linear = [x.replace('2', '11') for x in no_linear] param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in no_linear)] #param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in missing_keys)] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] t_total = num_train_steps optimizer = BertAdam(optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=t_total) global_step = 0 if args.do_train: logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_data)) logger.info(" Batch size = %d", args.train_batch_size) logger.info(" Num steps = %d", num_train_steps) model.train() import datetime fout = open(os.path.join(args.output_dir, "loss.{}".format(datetime.datetime.now())), 'w') for _ in trange(int(args.num_train_epochs), desc="Epoch"): tr_loss = 0 nb_tr_examples, nb_tr_steps = 0, 0 for step, batch in enumerate(tqdm(train_iterator.next_epoch_itr(), desc="Iteration")): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels = batch loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels) if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps loss.backward() fout.write("{} {}\n".format(loss.item()*args.gradient_accumulation_steps, original_loss.item())) tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 if (step + 1) % args.gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses lr_this_step = args.learning_rate * warmup_linear(global_step/t_total, args.warmup_proportion) for param_group in optimizer.param_groups: param_group['lr'] = lr_this_step optimizer.step() optimizer.zero_grad() global_step += 1 fout.close() # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") torch.save(model_to_save.state_dict(), output_model_file)
def main(): parser = argparse.ArgumentParser() ## Required parameters parser.add_argument("--data_dir", default=None, type=str, required=True, help="The input data dir. Should contain the .tsv files (or other data files) for the task.") parser.add_argument("--bert_model", default=None, type=str, required=True, help="Bert pre-trained model selected in the list: bert-base-uncased, " "bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.") parser.add_argument("--task_name", default=None, type=str, required=True, help="The name of the task to train.") parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.") ## Other parameters parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after WordPiece tokenization. \n" "Sequences longer than this will be truncated, and sequences shorter \n" "than this will be padded.") parser.add_argument("--do_train", default=False, action='store_true', help="Whether to run training.") parser.add_argument("--do_eval", default=False, action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--do_lower_case", default=False, action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=32, type=int, help="Total batch size for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Total batch size for eval.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--num_train_epochs", #default=3.0, default=1.0, type=float, help="Total number of training epochs to perform.") parser.add_argument("--warmup_proportion", default=0.1, type=float, help="Proportion of training to perform linear learning rate warmup for. " "E.g., 0.1 = 10%% of training.") parser.add_argument("--no_cuda", default=False, action='store_true', help="Whether not to use CUDA when available") parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument('--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument('--fp16', default=False, action='store_true', help="Whether to use 16-bit float precision instead of 32-bit") parser.add_argument('--loss_scale', type=float, default=0, help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n" "0 (default value): dynamic loss scaling.\n" "Positive power of 2: static loss scaling value.\n") ##########ADD## parser.add_argument("--K_V_dim", type=int, default=100, help="Key and Value dim == KG representation dim") parser.add_argument("--Q_dim", type=int, default=768, help="Query dim == Bert six output layer representation dim") parser.add_argument('--graphsage', default=False, action='store_true', help="Whether to use Attention GraphSage instead of GAT") parser.add_argument('--self_att', default=True, action='store_true', help="Whether to use GAT") ############### args = parser.parse_args() if args.local_rank == -1 or args.no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format( device, n_gpu, bool(args.local_rank != -1), args.fp16)) if args.gradient_accumulation_steps < 1: raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format( args.gradient_accumulation_steps)) args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) if not args.do_train and not args.do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir)) os.makedirs(args.output_dir, exist_ok=True) task_name = args.task_name.lower() train_data = None num_train_steps = None if args.do_train: # TODO import indexed_dataset from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler, BatchSampler import iterators #train_data = indexed_dataset.IndexedCachedDataset(args.data_dir) train_data = indexed_dataset.IndexedDataset(args.data_dir, fix_lua_indexing=True) #print(train_data) #print("-----------") if args.local_rank == -1: train_sampler = RandomSampler(train_data) else: train_sampler = DistributedSampler(train_data) train_sampler = BatchSampler(train_sampler, args.train_batch_size, True) def collate_fn(x): x = torch.LongTensor([xx for xx in x]) #x = torch.LongTensor([xx%9 for xx in x]) ## ### #entity_idx = x[:, 4*args.max_seq_length:5*args.max_seq_length] ### ### entity_idx = x[:, 3*args.max_seq_length:4*args.max_seq_length] ### #entity_idx = x[:, 4*args.max_seq_length] #print(entity_idx) #print(entity_idx.shape) # Build candidate # #print(entity_idx) uniq_idx = np.unique(entity_idx.numpy()) #print(uniq_idx) #print(uniq_idx.shape) #exit() #ent_candidate = embed(torch.LongTensor(uniq_idx+1)) #print(ent_candidate) #print(ent_candidate.shape) ent_candidate = torch.LongTensor(uniq_idx+1) #del #print(ent_candidate) #print(ent_candidate.shape) #print(ent_candidate) #print(ent_candidate.shape) #print(ent_candidate) #print(ent_candidate.shape) #print("================") ent_candidate = ent_candidate.repeat([n_gpu, 1]) #batch #print(ent_candidate) #print(ent_candidate.shape) #exit() #print(ent_candidate) #print(ent_candidate.shape) #exit() #ent_candidate = embed(torch.LongTensor(uniq_idx+1)) #del #print(ent_candidate) #print(ent_candidate.shape) #print(ent_candidate.size()) #exit() # #ent_candidate = torch.LongTensor(ent_candidate+1) #single #! --> return uniq_idx =>all entity in batch # build entity labels d = {} dd = [] for i, idx in enumerate(uniq_idx): d[idx] = i dd.append(idx) ### ''' ent_size = len(uniq_idx)-1 def map(x): if x == -1: return -1 else: rnd = random.uniform(0, 1) if rnd < 0.05: return dd[random.randint(1, ent_size)] elif rnd < 0.2: return -1 else: return x ''' ### ent_labels = entity_idx.clone() d[-1] = -1 ent_labels = ent_labels.apply_(lambda x: d[x]) ### ''' entity_idx.apply_(map) #ent_emb = embed(entity_idx+1) ent_emb = entity_idx+1 ## #! --> return entity+1 => input_ent mask = entity_idx.clone() mask.apply_(lambda x: 0 if x == -1 else 1) mask[:,0] = 1 ''' ### ### # entity_idx.apply_(map) #mask = entity_idx.clone() #mask.apply_(lambda x: 0 if x == -1 else 1) mask = x[:, 4*args.max_seq_length:5*args.max_seq_length] mask[:,0] = 1 entity_idx = entity_idx * mask ### entity_idx[entity_idx == 0] = -1 ### ent_emb = entity_idx+1 ### ### #return x[:,:args.max_seq_length], x[:,args.max_seq_length:2*args.max_seq_length], x[:,2*args.max_seq_length:3*args.max_seq_length], x[:,3*args.max_seq_length:4*args.max_seq_length], ent_emb, mask, x[:,6*args.max_seq_length:], ent_candidate, ent_labels ### ### return x[:,:args.max_seq_length], x[:,args.max_seq_length:2*args.max_seq_length], x[:,2*args.max_seq_length:3*args.max_seq_length], ent_emb, mask, x[:,5*args.max_seq_length:], ent_candidate, ent_labels ### train_iterator = iterators.EpochBatchIterator(train_data, collate_fn, train_sampler) num_train_steps = int( len(train_data) / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs) # Prepare model #model, missing_keys = BertForPreTraining.from_pretrained(args.bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank)) model, missing_keys = BertForPreTraining.from_pretrained(args.bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank), args=args) if args.fp16: model.half() model.to(device) if args.local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") model = DDP(model) elif n_gpu > 1: model = torch.nn.DataParallel(model) # Prepare optimizer param_optimizer = list(model.named_parameters()) no_linear = ['layer.2.output.dense_ent', 'layer.2.intermediate.dense_1', 'bert.encoder.layer.2.intermediate.dense_1_ent', 'layer.2.output.LayerNorm_ent'] no_linear = [x.replace('2', '11') for x in no_linear] param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in no_linear)] #param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in missing_keys)] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] t_total = num_train_steps if args.local_rank != -1: t_total = t_total // torch.distributed.get_world_size() if args.fp16: try: from apex.optimizers import FP16_Optimizer from apex.optimizers import FusedAdam except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") optimizer = FusedAdam(optimizer_grouped_parameters, lr=args.learning_rate, bias_correction=False, max_grad_norm=1.0) if args.loss_scale == 0: optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True) else: optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale) #logger.info(dir(optimizer)) #op_path = os.path.join(args.bert_model, "pytorch_op.bin") #optimizer.load_state_dict(torch.load(op_path)) else: optimizer = BertAdam(optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=t_total) global_step = 0 if args.do_train: logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_data)) logger.info(" Batch size = %d", args.train_batch_size) logger.info(" Num steps = %d", num_train_steps) model.train() import datetime fout = open(os.path.join(args.output_dir, "loss.{}".format(datetime.datetime.now())), 'w') more_than_one_2 = 0 less_than_one_2 = 0 for _ in trange(int(args.num_train_epochs), desc="Epoch"): tr_loss = 0 nb_tr_examples, nb_tr_steps = 0, 0 ### ''' for step, batch in enumerate(tqdm(train_iterator.next_epoch_itr(), desc="Iteration")): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels = batch ''' ### ### if args.local_rank == 0 or args.local_rank == -1: iters = tqdm(train_iterator.next_epoch_itr(), desc="Iteration") else: iters = train_iterator.next_epoch_itr() for step, batch in enumerate(iters): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels = batch ### ### #print(len(input_ids[input_ids==2])) if len(input_ids[input_ids==2]) != args.train_batch_size: for i_th_1, input_id in enumerate(input_ids): print(input_id[input_id==2]) print(len(input_id[input_id==2])) if len(input_id[input_id==2]) > 1: for i_th_2 ,id in enumerate(input_id): if id == 2: print("Befor:",input_id) input_ids[i_th_1][i_th_2] = 0 more_than_one_2 += 1 print("more_than_one_2:",more_than_one_2) print("After:",input_id) if len(input_id[input_id==2] == 1): break elif len(input_id[input_id==2]) < 1: print("Error!! Have no id=2 </s>") less_than_one_2 += 1 print("less_than_one_2:",less_than_one_2) print(input_id) input_ids[i_th_1][-1] = 2 else: print("ids_2 == 1") ### ### #start_time_1 = time.time() k_1, v_1, k_2, v_2, k_cand_1, v_cand_1, k_cand_2, v_cand_2, cand_pos_tensor = load_k_v_queryR_small(input_ent,ent_candidate) #k, v = load_k_v_queryR(input_ent,device) #input_ent_neighbor_emb, input_ent_r_emb, input_ent_outORin_emb = load_k_v_queryR(input_ent) #end_time_1 = time.time() #print("load_k_v_queryR:{}".format(end_time_1-start_time_1)) #print(ent_candidate) #print(ent_candidate.shape) #exit() #k_cand, v_cand = load_k_v_queryR_small(ent_candidate,"candidate") #k_cand, v_cand = load_k_v_queryR(ent_candidate,device) #input_ent_neighbor_emb_c, input_ent_r_emb_c, input_ent_outORin_emb_c = load_k_v_queryR(candidate) #end_time_2 = time.time() #print("load_cand:{}".format(end_time_2-end_time_1)) #k, v = load_batch_k_v_queryE(input_ent,500) #k_cand, v_cand = load_batch_k_v_queryE(ent_candidate,500) #k, v = load_batch_k_v_queryR(input_ent,300) #k_cand, v_cand = load_batch_k_v_queryR(ent_candidate,300) if args.fp16: #loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels, k_1.half(), v_1.half(), k_2.half(), v_2.half(), k_cand_1.half(), v_cand_1.half(), k_cand_2.half(), v_cand_2.half(), cand_pos_tensor) loss, original_loss = model(input_ids, None, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels, k_1.half(), v_1.half(), k_2.half(), v_2.half(), k_cand_1.half(), v_cand_1.half(), k_cand_2.half(), v_cand_2.half(), cand_pos_tensor) else: #loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels, k_1, v_1, k_2, v_2, k_cand_1, v_cand_1, k_cand_2, v_cand_2, cand_pos_tensor) loss, original_loss = model(input_ids, None, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels, k_1, v_1, k_2, v_2, k_cand_1, v_cand_1, k_cand_2, v_cand_2, cand_pos_tensor) if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. original_loss = original_loss.mean() if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps if args.fp16: optimizer.backward(loss) else: loss.backward() end_time_4 = time.time() #print("bp time:{}".format(end_time_4)) #print("=====================================") fout.write("{} {}\n".format(loss.item()*args.gradient_accumulation_steps, original_loss.item())) tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 if (step + 1) % args.gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses lr_this_step = args.learning_rate * warmup_linear(global_step/t_total, args.warmup_proportion) for param_group in optimizer.param_groups: param_group['lr'] = lr_this_step optimizer.step() optimizer.zero_grad() global_step += 1 if global_step % 100000 == 0: model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin_{}".format(global_step)) torch.save(model_to_save.state_dict(), output_model_file) fout.close() # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") torch.save(model_to_save.state_dict(), output_model_file)
def main(): parser = argparse.ArgumentParser() ## Required parameters parser.add_argument("--data_dir", default=None, type=str, required=True, help="The input data dir. Should contain the .tsv files (or other data files) for the task.") parser.add_argument("--bert_model", default=None, type=str, required=True, help="Bert pre-trained model selected in the list: bert-base-uncased, " "bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.") parser.add_argument("--task_name", default=None, type=str, required=True, help="The name of the task to train.") parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.") ## Other parameters parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after WordPiece tokenization. \n" "Sequences longer than this will be truncated, and sequences shorter \n" "than this will be padded.") parser.add_argument("--do_train", default=False, action='store_true', help="Whether to run training.") parser.add_argument("--do_eval", default=False, action='store_true', help="Whether to run eval on the dev set.") parser.add_argument("--do_lower_case", default=False, action='store_true', help="Set this flag if you are using an uncased model.") parser.add_argument("--train_batch_size", default=32, type=int, help="Total batch size for training.") parser.add_argument("--eval_batch_size", default=8, type=int, help="Total batch size for eval.") parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.") parser.add_argument("--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.") parser.add_argument("--warmup_proportion", default=0.1, type=float, help="Proportion of training to perform linear learning rate warmup for. " "E.g., 0.1 = 10%% of training.") parser.add_argument("--no_cuda", default=False, action='store_true', help="Whether not to use CUDA when available") parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus") parser.add_argument('--seed', type=int, default=42, help="random seed for initialization") parser.add_argument('--gradient_accumulation_steps', type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.") parser.add_argument('--fp16', default=False, action='store_true', help="Whether to use 16-bit float precision instead of 32-bit") parser.add_argument('--loss_scale', type=float, default=0, help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n" "0 (default value): dynamic loss scaling.\n" "Positive power of 2: static loss scaling value.\n") args = parser.parse_args() if args.local_rank == -1 or args.no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(args.local_rank) device = torch.device("cuda", args.local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format( device, n_gpu, bool(args.local_rank != -1), args.fp16)) if args.gradient_accumulation_steps < 1: raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format( args.gradient_accumulation_steps)) args.train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if n_gpu > 0: torch.cuda.manual_seed_all(args.seed) if not args.do_train and not args.do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") if os.path.exists(args.output_dir) and os.listdir(args.output_dir): raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir)) os.makedirs(args.output_dir, exist_ok=True) task_name = args.task_name.lower() vecs = [] vecs.append([0]*200) # 扩充CLS的位置,其他所有索引向后+1. with open("config_data/kg_embed/entity2vec.vec", 'r') as fin: #with open("pretrain_data/config_data/entity2vec.vec", 'r') as fin: for line in fin: vec = line.strip().split('\t') #vec = [float(x) for x in vec if x != ""] vec = [float(x) for x in vec] vecs.append(vec) print("vecs_len=%s" % str(len(vecs))) print("vecs_dim=%s" % str(len(vecs[0]))) ent_embed = torch.FloatTensor(vecs) ent_embed = torch.nn.Embedding.from_pretrained(ent_embed) #ent_embed = torch.nn.Embedding(5041175, 100) logger.info("Shape of entity embedding: "+str(ent_embed.weight.size())) vecs = [] vecs.append([0] * 4096) # 扩充CLS的位置,其他所有索引向后+1. with open("config_data/kg_embed/image2vec.vec", 'r') as fin: #with open("pretrain_data/image_vec/image2vec.vec", 'r') as fin: for line in fin: vec = line.strip().split('\t') vec = [float(x) for x in vec] vecs.append(vec) print("vecs_len=%s" % str(len(vecs))) print("vecs_dim=%s" % str(len(vecs[0]))) img_embed = torch.FloatTensor(vecs) img_embed = torch.nn.Embedding.from_pretrained(img_embed) logger.info("Shape of image embedding: " + str(img_embed.weight.size())) del vecs train_data = None num_train_steps = None if args.do_train: # TODO import indexed_dataset from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler,BatchSampler import iterators #train_data = indexed_dataset.IndexedCachedDataset(args.data_dir) train_data = indexed_dataset.IndexedDataset(args.data_dir, fix_lua_indexing=True) if args.local_rank == -1: train_sampler = RandomSampler(train_data) else: train_sampler = DistributedSampler(train_data) train_sampler = BatchSampler(train_sampler, args.train_batch_size, True) def collate_fn(x): x = torch.LongTensor([xx for xx in x]) entity_idx = x[:, 4 * args.max_seq_length:5 * args.max_seq_length] print("entity_idx=%s" % entity_idx) image_idx = x[:, 6 * args.max_seq_length:7 * args.max_seq_length] print("image_idx=%s" % image_idx) # Build candidate ent_uniq_idx = np.unique(entity_idx.numpy()) print("ent_uniq_idx=%s" % str(ent_uniq_idx)) img_uniq_idx = np.unique(image_idx.numpy()) print("img_uniq_idx=%s" % str(img_uniq_idx)) ent_candidate = ent_embed(torch.LongTensor(ent_uniq_idx + 1)) ent_candidate = ent_candidate.repeat([n_gpu, 1]) img_candidate = img_embed(torch.LongTensor(img_uniq_idx + 1)) img_candidate = img_candidate.repeat([n_gpu, 1]) # build entity labels ent_idx_dict = {} ent_idx_list = [] for idx, idx_value in enumerate(ent_uniq_idx): ent_idx_dict[idx_value] = idx ent_idx_list.append(idx_value) ent_size = len(ent_uniq_idx)-1 # build image labels img_idx_dict = {} img_idx_list = [] for idx, idx_value in enumerate(img_uniq_idx): img_idx_dict[idx_value] = idx img_idx_list.append(idx_value) img_size = len(img_uniq_idx) - 1 def ent_map(x): if x == -1: return -1 else: rnd = random.uniform(0, 1) if rnd < 0.05: return ent_idx_list[random.randint(1, ent_size)] elif rnd < 0.2: return -1 else: return x def img_map(x): if x == -1: return -1 else: rnd = random.uniform(0, 1) if rnd < 0.05: return img_idx_list[random.randint(1, ent_size)] elif rnd < 0.2: return -1 else: return x ent_labels = entity_idx.clone() ent_idx_dict[-1] = -1 ent_labels = ent_labels.apply_(lambda x: ent_idx_dict[x]) entity_idx.apply_(ent_map) ent_emb = ent_embed(entity_idx+1) ent_mask = entity_idx.clone() ent_mask.apply_(lambda x: 0 if x == -1 else 1) ent_mask[:,0] = 1 img_labels = image_idx.clone() img_idx_dict[-1] = -1 img_labels = img_labels.apply_(lambda x: img_idx_dict[x]) image_idx.apply_(img_map) img_emb = img_embed(image_idx + 1) img_mask = image_idx.clone() img_mask.apply_(lambda x: 0 if x == -1 else 1) img_mask[:, 0] = 1 input_ids = x[:,:args.max_seq_length] input_mask = x[:,args.max_seq_length:2*args.max_seq_length] segment_ids = x[:,2*args.max_seq_length:3*args.max_seq_length] masked_lm_labels = x[:,3*args.max_seq_length:4*args.max_seq_length] next_sentence_label = x[:,8*args.max_seq_length:] return input_ids, input_mask, segment_ids, masked_lm_labels, ent_emb, ent_mask, img_emb, img_mask, next_sentence_label, ent_candidate, ent_labels, img_candidate, img_labels train_iterator = iterators.EpochBatchIterator(train_data, collate_fn, train_sampler) num_train_steps = int( len(train_data) / args.train_batch_size / args.gradient_accumulation_steps * args.num_train_epochs) print ("len(train_data)=%s" % len(train_data)) # Prepare model model, missing_keys = BertForPreTraining.from_pretrained(args.bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank)) if args.fp16: model.half() model.to(device) if args.local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") model = DDP(model) elif n_gpu > 1: model = torch.nn.DataParallel(model) # Prepare optimizer param_optimizer = list(model.named_parameters()) print ("param_optimizer:") #for param in model.named_parameters(): # print(param[0]) #no_linear = ['layer.2.output.dense_ent', 'layer.2.intermediate.dense_1', 'bert.encoder.layer.2.intermediate.dense_1_ent', 'layer.2.output.LayerNorm_ent'] #no_linear = [x.replace('2', '11') for x in no_linear] no_linear = ['layer.11.output.dense_entity', 'layer.11.output.LayerNorm_entity', 'layer.11.output.dense_image', 'layer.11.output.LayerNorm_entity'] param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in no_linear)] print ("param_optimizer--no_linear") #for param in param_optimizer: # print (param[0]) #param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in missing_keys)] #no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] #no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_token.bias', 'LayerNorm_token.weight', 'LayerNorm_entity.bias', 'LayerNorm_entity.weight', 'LayerNorm_image.bias', 'LayerNorm_image.weight'] optimizer_grouped_parameters = [ # weight decay to avoid overfitting # source: https://blog.csdn.net/program_developer/article/details/80867468 # source: https://blog.csdn.net/m0_37531129/article/details/101390592 {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, # the decay of bias and normalization.weight has nothing to do with weight decay {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] # optimizer_grouped_parameters_display is only used to debug # optimizer_grouped_parameters_display = [ # {'params': [(n,p) for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, # {'params': [(n,p) for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} # ] # print ("optimizer_grouped_parameters_display-0:") # for param in optimizer_grouped_parameters_display[0]['params']: # print (param[0]) # # print ("optimizer_grouped_parameters_display-1:") # for param in optimizer_grouped_parameters_display[1]['params']: # print (param[0]) t_total = num_train_steps if args.local_rank != -1: t_total = t_total // torch.distributed.get_world_size() if args.fp16: try: #from apex.optimizers import FP16_Optimizer from apex.fp16_utils.fp16_optimizer import FP16_Optimizer from apex.optimizers import FusedAdam except ImportError: raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") #optimizer = FusedAdam(optimizer_grouped_parameters, # lr=args.learning_rate, # bias_correction=False, # max_grad_norm=1.0) optimizer = FusedAdam(optimizer_grouped_parameters, lr=args.learning_rate, bias_correction=False) if args.loss_scale == 0: optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True) else: optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale) #logger.info(dir(optimizer)) #op_path = os.path.join(args.bert_model, "pytorch_op.bin") #optimizer.load_state_dict(torch.load(op_path)) else: optimizer = BertAdam(optimizer_grouped_parameters, lr=args.learning_rate, warmup=args.warmup_proportion, t_total=t_total) global_step = 0 if args.do_train: logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_data)) logger.info(" Batch size = %d", args.train_batch_size) logger.info(" Num steps = %d", num_train_steps) model.train() import datetime fout = open(os.path.join(args.output_dir, "loss.{}".format(datetime.datetime.now())), 'w') for _ in trange(int(args.num_train_epochs), desc="Epoch"): tr_loss = 0 nb_tr_examples, nb_tr_steps = 0, 0 for step, batch in enumerate(tqdm(train_iterator.next_epoch_itr(), desc="Iteration")): print ("step=%s" % str(step)) print ("len(batch)=%s" % str(len(batch))) batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, masked_lm_labels, input_ent, ent_mask, input_img, img_mask, next_sentence_label, ent_candidate, ent_labels, img_candidate, img_labels = batch print ("\ninput_ids.size=%s" % str(input_ids.size())) print ("input_mask.size=%s" % str(input_mask.size())) print ("segment_ids.size=%s" % str(segment_ids.size())) print ("masked_lm_labels.size=%s" % str(masked_lm_labels.size())) print ("input_ent.size=%s" % str(input_ent.size())) print ("ent_mask.size=%s" % str(ent_mask.size())) print ("input_img.size=%s" % str(input_img.size())) print ("img_mask.size=%s" % str(img_mask.size())) print ("next_sentence_label.size=%s" % str(next_sentence_label.size())) print ("ent_candidate.size=%s" % str(ent_candidate.size())) print ("ent_labels.size=%s" % str(ent_labels.size())) print ("img_candidate.size=%s" % str(img_candidate.size())) print ("img_labels.size=%s" % str(img_labels.size())) if args.fp16: loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent.half(), ent_mask, input_img.half(), img_mask, next_sentence_label, ent_candidate.half(), ent_labels, img_candidate.half(), img_labels) else: loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, input_img, img_mask, next_sentence_label, ent_candidate, ent_labels, img_candidate, img_labels) if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. original_loss = original_loss.mean() if args.gradient_accumulation_steps > 1: loss = loss / args.gradient_accumulation_steps print("\nloss=%s\n" % str(loss)) if args.fp16: optimizer.backward(loss) else: loss.backward() fout.write("{} {}\n".format(loss.item()*args.gradient_accumulation_steps, original_loss.item())) tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 if (step + 1) % args.gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses # source: https://blog.csdn.net/m0_37531129/article/details/101390592 lr_this_step = args.learning_rate * warmup_linear(global_step/t_total, args.warmup_proportion) for param_group in optimizer.param_groups: param_group['lr'] = lr_this_step optimizer.step() optimizer.zero_grad() global_step += 1 #if global_step % 1000 == 0: # model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self # output_model_file = os.path.join(args.output_dir, "pytorch_model.bin_{}".format(global_step)) # torch.save(model_to_save.state_dict(), output_model_file) fout.close() # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") torch.save(model_to_save.state_dict(), output_model_file)
def main(): ## Required parameters path_to_ernie = "" data_dir = os.path.join(path_to_ernie, "pretrain_data/merge") bert_model = os.path.join(path_to_ernie, "ernie_base") task_name = "pretrain" output_dir = os.path.join(path_to_ernie, "pretrain_out") max_seq_length = 256 do_train = True do_eval = False do_lower_case = False train_batch_size = 4 eval_batch_size = 8 learning_rate = 5e-5 num_train_epochs = 3.0 warmup_proportion = default = 0.1 no_cuda = False local_rank = -1 seed = 42 gradient_accumulation_steps = 1 fp16 = True loss_scale = 0.0 if local_rank == -1 or no_cuda: device = torch.device("cuda" if torch.cuda.is_available() and not no_cuda else "cpu") n_gpu = torch.cuda.device_count() else: torch.cuda.set_device(local_rank) device = torch.device("cuda", local_rank) n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs torch.distributed.init_process_group(backend='nccl') logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format( device, n_gpu, bool(local_rank != -1), fp16)) if gradient_accumulation_steps < 1: raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format( gradient_accumulation_steps)) train_batch_size = int(train_batch_size / gradient_accumulation_steps) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if n_gpu > 0: torch.cuda.manual_seed_all(seed) if not do_train and not do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") if os.path.exists(output_dir) and os.listdir(output_dir): import shutil shutil.rmtree(output_dir) # raise ValueError("Output directory ({}) already exists and is not empty.".format(output_dir)) os.makedirs(output_dir, exist_ok=True) task_name = task_name.lower() vecs = [] vecs.append([0] * 100) # CLS with open(os.path.join(path_to_ernie, "kg_embed/entity2vec.vec"), 'r') as fin: for line in fin: # vec = line.strip().split('\t') vec = line.strip().split(' ') vec = [float(x) for x in vec] vecs.append(vec) embed = torch.FloatTensor(vecs) embed = torch.nn.Embedding.from_pretrained(embed) # embed = torch.nn.Embedding(5041175, 100) logger.info("Shape of entity embedding: " + str(embed.weight.size())) del vecs train_data = None num_train_steps = None if do_train: # TODO import indexed_dataset from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler, BatchSampler import iterators # train_data = indexed_dataset.IndexedCachedDataset(data_dir) train_data = indexed_dataset.IndexedDataset(data_dir, fix_lua_indexing=True) if local_rank == -1: train_sampler = RandomSampler(train_data) else: train_sampler = DistributedSampler(train_data) train_sampler = BatchSampler(train_sampler, train_batch_size, True) def collate_fn(x): x = torch.LongTensor([xx for xx in x]) entity_idx = x[:, 4 * max_seq_length:5 * max_seq_length] # Build candidate uniq_idx = np.unique(entity_idx.numpy()) ent_candidate = embed(torch.LongTensor(uniq_idx + 1)) ent_candidate = ent_candidate.repeat([n_gpu, 1]) # build entity labels d = {} dd = [] for i, idx in enumerate(uniq_idx): d[idx] = i dd.append(idx) ent_size = len(uniq_idx) - 1 def map(x): if x == -1: return -1 else: rnd = random.uniform(0, 1) if rnd < 0.05: return dd[random.randint(1, ent_size)] elif rnd < 0.2: return -1 else: return x ent_labels = entity_idx.clone() d[-1] = -1 ent_labels = ent_labels.apply_(lambda x: d[x]) entity_idx.apply_(map) ent_emb = embed(entity_idx + 1) mask = entity_idx.clone() mask.apply_(lambda x: 0 if x == -1 else 1) mask[:, 0] = 1 return x[:, :max_seq_length], x[:, max_seq_length:2 * max_seq_length], x[:, 2 * max_seq_length:3 * max_seq_length], x[:,3 * max_seq_length:4 * max_seq_length], ent_emb, mask, x[:,6 * max_seq_length:], ent_candidate, ent_labels train_iterator = iterators.EpochBatchIterator(train_data, collate_fn, train_sampler) num_train_steps = int( len(train_data) / train_batch_size / gradient_accumulation_steps * num_train_epochs) # Prepare model model, missing_keys = BertForPreTraining.from_pretrained(bert_model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format( local_rank)) # if fp16: # model.half() model.to(device) if local_rank != -1: try: from apex.parallel import DistributedDataParallel as DDP except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") model = DDP(model) elif n_gpu > 1: model = torch.nn.DataParallel(model) # Prepare optimizer param_optimizer = list(model.named_parameters()) no_linear = ['layer.2.output.dense_ent', 'layer.2.intermediate.dense_1', 'bert.encoder.layer.2.intermediate.dense_1_ent', 'layer.2.output.LayerNorm_ent'] no_linear = [x.replace('2', '11') for x in no_linear] param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in no_linear)] # param_optimizer = [(n, p) for n, p in param_optimizer if not any(nl in n for nl in missing_keys)] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm_ent.bias', 'LayerNorm_ent.weight'] optimizer_grouped_parameters = [ {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] t_total = num_train_steps if local_rank != -1: t_total = t_total // torch.distributed.get_world_size() if fp16: try: # from apex.optimizers import FP16_Optimizer # from apex.optimizers import FusedAdam # from apex.contrib.optimizers import FP16_Optimizer from apex.optimizers import FusedAdam import apex.amp as amp except ImportError: raise ImportError( "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.") # optimizer = FusedAdam(optimizer_grouped_parameters, # lr=learning_rate, # bias_correction=False, # max_grad_norm=1.0) optimizer = FusedAdam(optimizer_grouped_parameters, lr=learning_rate, bias_correction=False) model, optimizer = amp.initialize(model, optimizer, opt_level="O3") # if loss_scale == 0: # optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True) # else: # optimizer = FP16_Optimizer(optimizer, static_loss_scale=loss_scale) # logger.info(dir(optimizer)) # op_path = os.path.join(bert_model, "pytorch_op.bin") # optimizer.load_state_dict(torch.load(op_path)) else: optimizer = BertAdam(optimizer_grouped_parameters, lr=learning_rate, warmup=warmup_proportion, t_total=t_total) global_step = 0 if do_train: logger.info("***** Running training *****") logger.info(" Num examples = %d", len(train_data)) logger.info(" Batch size = %d", train_batch_size) logger.info(" Num steps = %d", num_train_steps) model.train() import datetime fout = open(os.path.join(output_dir, "loss.{}".format(datetime.datetime.now())), 'w') for _ in trange(int(num_train_epochs), desc="Epoch"): tr_loss = 0 nb_tr_examples, nb_tr_steps = 0, 0 for step, batch in enumerate(tqdm(train_iterator.next_epoch_itr(), desc="Iteration")): batch = tuple(t.to(device) for t in batch) input_ids, input_mask, segment_ids, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels = batch if fp16: loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent.half(), ent_mask, next_sentence_label, ent_candidate.half(), ent_labels) else: loss, original_loss = model(input_ids, segment_ids, input_mask, masked_lm_labels, input_ent, ent_mask, next_sentence_label, ent_candidate, ent_labels) if n_gpu > 1: loss = loss.mean() # mean() to average on multi-gpu. original_loss = original_loss.mean() if gradient_accumulation_steps > 1: loss = loss / gradient_accumulation_steps if fp16: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() # optimizer.backward(loss) else: loss.backward() fout.write("{} {}\n".format(loss.item() * gradient_accumulation_steps, original_loss.item())) tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 if (step + 1) % gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses lr_this_step = learning_rate * warmup_linear(global_step / t_total, warmup_proportion) for param_group in optimizer.param_groups: param_group['lr'] = lr_this_step optimizer.step() optimizer.zero_grad() global_step += 1 # if global_step % 1000 == 0: # model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self # output_model_file = os.path.join(output_dir, "pytorch_model.bin_{}".format(global_step)) # torch.save(model_to_save.state_dict(), output_model_file) fout.close() # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(output_dir, "pytorch_model.bin") torch.save(model_to_save.state_dict(), output_model_file)