from data_util import UMLSDataset, fixed_length_dataloader
from model import UMLSPretrainedModel
from transformers import AdamW, get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup, get_constant_schedule_with_warmup
from tqdm import tqdm, trange
import torch
from torch import nn
import time
import os
import numpy as np
import argparse
import time
import pathlib
#import ipdb
# try:
# from torch.utils.tensorboard import SummaryWriter
# except:
from tensorboardX import SummaryWriter
def train(args, model, train_dataloader, umls_dataset):
writer = SummaryWriter(comment='umls')
t_total = args.max_steps
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
{
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
"weight_decay": args.weight_decay,
},
{"params": [p for n, p in model.named_parameters() if any(
nd in n for nd in no_decay)], "weight_decay": 0.0},
]
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate, eps=args.adam_epsilon)
args.warmup_steps = int(args.warmup_steps)
if args.schedule == 'linear':
scheduler = get_linear_schedule_with_warmup(
optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
)
if args.schedule == 'constant':
scheduler = get_constant_schedule_with_warmup(
optimizer, num_warmup_steps=args.warmup_steps
)
if args.schedule == 'cosine':
scheduler = get_cosine_schedule_with_warmup(
optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
)
print("***** Running training *****")
print(" Total Steps =", t_total)
print(" Steps needs to be trained=", t_total - args.shift)
print(" Instantaneous batch size per GPU =", args.train_batch_size)
print(
" Total train batch size (w. parallel, distributed & accumulation) =",
args.train_batch_size
* args.gradient_accumulation_steps,
)
print(" Gradient Accumulation steps =", args.gradient_accumulation_steps)
model.zero_grad()
for i in range(args.shift):
scheduler.step()
global_step = args.shift
while True:
model.train()
epoch_iterator = tqdm(train_dataloader, desc="Iteration", ascii=True)
batch_loss = 0.
batch_sty_loss = 0.
batch_cui_loss = 0.
batch_re_loss = 0.
for _, batch in enumerate(epoch_iterator):
input_ids_0 = batch[0].to(args.device)
input_ids_1 = batch[1].to(args.device)
input_ids_2 = batch[2].to(args.device)
cui_label_0 = batch[3].to(args.device)
cui_label_1 = batch[4].to(args.device)
cui_label_2 = batch[5].to(args.device)
sty_label_0 = batch[6].to(args.device)
sty_label_1 = batch[7].to(args.device)
sty_label_2 = batch[8].to(args.device)
# use batch[9] for re, use batch[10] for rel
if args.use_re:
re_label = batch[9].to(args.device)
else:
re_label = batch[10].to(args.device)
# for item in batch:
# print(item.shape)
loss, (sty_loss, cui_loss, re_loss) = \
model(input_ids_0, input_ids_1, input_ids_2,
cui_label_0, cui_label_1, cui_label_2,
sty_label_0, sty_label_1, sty_label_2,
re_label)
batch_loss = float(loss.item())
batch_sty_loss = float(sty_loss.item())
batch_cui_loss = float(cui_loss.item())
batch_re_loss = float(re_loss.item())
# tensorboardX
writer.add_scalar(
'rel_count', train_dataloader.batch_sampler.rel_sampler_count, global_step=global_step)
writer.add_scalar('batch_loss', batch_loss,
global_step=global_step)
writer.add_scalar('batch_sty_loss', batch_sty_loss,
global_step=global_step)
writer.add_scalar('batch_cui_loss', batch_cui_loss,
global_step=global_step)
writer.add_scalar('batch_re_loss', batch_re_loss,
global_step=global_step)
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
loss.backward()
epoch_iterator.set_description("Rel_count: %s, Loss: %0.4f, Sty: %0.4f, Cui: %0.4f, Re: %0.4f" %
(train_dataloader.batch_sampler.rel_sampler_count, batch_loss, batch_sty_loss, batch_cui_loss, batch_re_loss))
if (global_step + 1) % args.gradient_accumulation_steps == 0:
torch.nn.utils.clip_grad_norm_(
model.parameters(), args.max_grad_norm)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if global_step % args.save_step == 0 and global_step > 0:
save_path = os.path.join(
args.output_dir, f'model_{global_step}.pth')
torch.save(model, save_path)
# re_embedding
if args.use_re:
writer.add_embedding(model.re_embedding.weight, metadata=umls_dataset.re2id.keys(
), global_step=global_step, tag="re embedding")
else:
# print(len(umls_dataset.rel2id))
# print(model.re_embedding.weight.shape)
writer.add_embedding(model.re_embedding.weight, metadata=umls_dataset.rel2id.keys(
), global_step=global_step, tag="rel embedding")
# sty_parameter
writer.add_embedding(model.linear_sty.weight, metadata=umls_dataset.sty2id.keys(
), global_step=global_step, tag="sty weight")
if args.max_steps > 0 and global_step > args.max_steps:
return None
return None
def run(args):
torch.manual_seed(args.seed) # cpu
torch.cuda.manual_seed(args.seed) # gpu
np.random.seed(args.seed) # numpy
torch.backends.cudnn.deterministic = True # cudnn
#args.output_dir = args.output_dir + "_" + str(int(time.time()))
# dataloader
if args.lang == "eng":
lang = ["ENG"]
if args.lang == "all":
lang = None
assert args.model_name_or_path.find("bio") == -1, "Should use multi-language model"
umls_dataset = UMLSDataset(
umls_folder=args.umls_dir, model_name_or_path=args.model_name_or_path, lang=lang, json_save_path=args.output_dir)
umls_dataloader = fixed_length_dataloader(
umls_dataset, fixed_length=args.train_batch_size, num_workers=args.num_workers)
if args.use_re:
rel_label_count = len(umls_dataset.re2id)
else:
rel_label_count = len(umls_dataset.rel2id)
model_load = False
if os.path.exists(args.output_dir):
save_list = []
for f in os.listdir(args.output_dir):
if f[0:5] == "model" and f[-4:] == ".pth":
save_list.append(int(f[6:-4]))
if len(save_list) > 0:
args.shift = max(save_list)
if os.path.exists(os.path.join(args.output_dir, 'last_model.pth')):
model = torch.load(os.path.join(
args.output_dir, 'last_model.pth')).to(args.device)
model_load = True
else:
model = torch.load(os.path.join(
args.output_dir, f'model_{max(save_list)}.pth')).to(args.device)
model_load = True
if not model_load:
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
model = UMLSPretrainedModel(device=args.device, model_name_or_path=args.model_name_or_path,
cui_label_count=len(umls_dataset.cui2id),
rel_label_count=rel_label_count,
sty_label_count=len(umls_dataset.sty2id),
re_weight=args.re_weight,
sty_weight=args.sty_weight).to(args.device)
args.shift = 0
model_load = True
if args.do_train:
torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
train(args, model, umls_dataloader, umls_dataset)
torch.save(model, os.path.join(args.output_dir, 'last_model.pth'))
return None
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--umls_dir",
default="../umls",
type=str,
help="UMLS dir",
)
parser.add_argument(
"--model_name_or_path",
default="../biobert_v1.1",
type=str,
help="Path to pre-trained model or shortcut name selected in the list: ",
)
parser.add_argument(
"--output_dir",
default="output",
type=str,
help="The output directory where the model predictions and checkpoints will be written.",
)
parser.add_argument(
"--save_step",
default=25000,
type=int,
help="Save step",
)
# Other parameters
parser.add_argument(
"--max_seq_length",
default=32,
type=int,
help="The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.",
)
parser.add_argument("--do_train", default=True, type=bool, help="Whether to run training.")
parser.add_argument(
"--train_batch_size", default=256, type=int, help="Batch size per GPU/CPU for training.",
)
parser.add_argument(
"--gradient_accumulation_steps",
type=int,
default=8,
help="Number of updates steps to accumulate before performing a backward/update pass.",
)
parser.add_argument("--learning_rate", default=2e-5,
type=float, help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay", default=0.01,
type=float, help="Weight decay if we apply some.")
parser.add_argument("--adam_epsilon", default=1e-8,
type=float, help="Epsilon for Adam optimizer.")
parser.add_argument("--max_grad_norm", default=1.0,
type=float, help="Max gradient norm.")
parser.add_argument(
"--max_steps",
default=1000000,
type=int,
help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
)
parser.add_argument("--warmup_steps", default=10000,
help="Linear warmup over warmup_steps or a float.")
parser.add_argument("--device", type=str, default='cuda:1', help="device")
parser.add_argument("--seed", type=int, default=72,
help="random seed for initialization")
parser.add_argument("--schedule", type=str, default="linear",
choices=["linear", "cosine", "constant"], help="Schedule.")
parser.add_argument("--trans_margin", type=float, default=1.0,
help="Margin of TransE.")
parser.add_argument("--use_re", default=False, type=bool,
help="Whether to use re or rel.")
parser.add_argument("--num_workers", default=1, type=int,
help="Num workers for data loader, only 0 can be used for Windows")
parser.add_argument("--lang", default='eng', type=str, choices=["eng", "all"],
help="language range, eng or all")
parser.add_argument("--sty_weight", type=float, default=0.0,
help="Weight of sty.")
parser.add_argument("--re_weight", type=float, default=1.0,
help="Weight of re.")
args = parser.parse_args()
run(args)
if __name__ == "__main__":
main()
