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)
