in scripts/train_qa.py [0:0]
def main():
args = train_args()
if args.fp16:
import apex
apex.amp.register_half_function(torch, 'einsum')
date_curr = date.today().strftime("%m-%d-%Y")
model_name = f"{args.prefix}-seed{args.seed}-bsz{args.train_batch_size}-fp16{args.fp16}-lr{args.learning_rate}-decay{args.weight_decay}-neg{args.neg_num}-sn{args.shared_norm}-adam{args.use_adam}-warm{args.warmup_ratio}-sp{args.sp_weight}"
args.output_dir = os.path.join(args.output_dir, date_curr, model_name)
tb_logger = SummaryWriter(os.path.join(args.output_dir.replace("logs","tflogs")))
if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
print(
f"output directory {args.output_dir} already exists and is not empty.")
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir, exist_ok=True)
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO,
handlers=[logging.FileHandler(os.path.join(args.output_dir, "log.txt")),
logging.StreamHandler()])
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.info(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:
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 %s n_gpu %d distributed training %r",
device, n_gpu, bool(args.local_rank != -1))
if args.shared_norm:
# chains of each question are on the same gpu
assert (args.train_batch_size // n_gpu) == args.neg_num + 1
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)
# define model
if args.model_name == "spanbert":
bert_config = AutoConfig.from_pretrained("/private/home/span-bert")
tokenizer = AutoTokenizer.from_pretrained('bert-large-cased')
else:
bert_config = AutoConfig.from_pretrained(args.model_name)
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
model = QAModel(bert_config, args)
collate_fc = partial(qa_collate, pad_id=tokenizer.pad_token_id)
eval_dataset = QADataset(tokenizer, args.predict_file, args.max_seq_len, args.max_q_len)
eval_dataloader = DataLoader(eval_dataset, batch_size=args.predict_batch_size, collate_fn=collate_fc, pin_memory=True, num_workers=args.num_workers)
logger.info(f"Num of dev batches: {len(eval_dataloader)}")
if args.init_checkpoint != "":
logger.info(f"Loading model from {args.init_checkpoint}")
model = load_saved(model, args.init_checkpoint)
model.to(device)
print(f"number of trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
if args.do_train:
no_decay = ['bias', 'LayerNorm.weight']
optimizer_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}
]
if args.use_adam:
optimizer = Adam(optimizer_parameters,
lr=args.learning_rate, eps=args.adam_epsilon)
else:
optimizer = AdamW(optimizer_parameters,
lr=args.learning_rate, eps=args.adam_epsilon)
if args.fp16:
from apex import amp
model, optimizer = amp.initialize(
model, optimizer, opt_level=args.fp16_opt_level)
else:
if args.fp16:
from apex import amp
model = amp.initialize(model, opt_level=args.fp16_opt_level)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
output_device=args.local_rank)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
if args.do_train:
global_step = 0 # gradient update step
batch_step = 0 # forward batch count
best_em = 0
train_loss_meter = AverageMeter()
model.train()
train_dataset = QADataset(tokenizer, args.train_file, args.max_seq_len, args.max_q_len, train=True)
train_sampler = MhopSampler(train_dataset, num_neg=args.neg_num)
train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size, pin_memory=True, collate_fn=collate_fc, num_workers=args.num_workers, sampler=train_sampler)
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
warmup_steps = t_total * args.warmup_ratio
scheduler = get_linear_schedule_with_warmup(
optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
)
logger.info('Start training....')
for epoch in range(int(args.num_train_epochs)):
for batch in tqdm(train_dataloader):
batch_step += 1
batch_inputs = move_to_cuda(batch["net_inputs"])
loss = model(batch_inputs)
if n_gpu > 1:
loss = loss.mean()
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
train_loss_meter.update(loss.item())
if (batch_step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(
model.parameters(), args.max_grad_norm)
optimizer.step()
scheduler.step()
model.zero_grad()
global_step += 1
# logger.info(f"current batch loss: {loss.item()}")
tb_logger.add_scalar('batch_train_loss',
loss.item(), global_step)
tb_logger.add_scalar('smoothed_train_loss',
train_loss_meter.avg, global_step)
if args.eval_period != -1 and global_step % args.eval_period == 0:
metrics = predict(args, model, eval_dataloader, logger)
em = metrics["em"]
logger.info("Step %d Train loss %.2f em %.2f on epoch=%d" % (global_step, train_loss_meter.avg, em*100, epoch))
if best_em < em:
logger.info("Saving model with best em %.2f -> em %.2f on step=%d" %
(best_em*100, em*100, global_step))
torch.save(model.state_dict(), os.path.join(
args.output_dir, f"checkpoint_best.pt"))
model = model.to(device)
best_em = em
metrics = predict(args, model, eval_dataloader, logger)
em = metrics["em"]
logger.info("Step %d Train loss %.2f em %.2f" % (
global_step, train_loss_meter.avg, em*100))
tb_logger.add_scalar('dev_em', em*100, global_step)
if best_em < em:
logger.info("Saving model with best em %.2f -> em %.2f on epoch=%d" % (best_em*100, em*100, epoch))
torch.save(model.state_dict(), os.path.join(
args.output_dir, f"checkpoint_best.pt"))
best_em = em
logger.info("Training finished!")
elif args.do_predict:
metrics = predict(args, model, eval_dataloader, logger, fixed_thresh=0.8)
logger.info(f"test performance {metrics}")
elif args.do_test:
eval_final(args, model, eval_dataloader, weight=0.8)