def get_latest_checkpoint(args):
if not os.path.exists(args.training_dir):
return args.init_model_dir, 0
files = list(next(os.walk(args.training_dir))[2])
def valid_checkpoint(checkpoint):
return checkpoint.startswith("checkpoint-")
logger.info("checkpoint files")
logger.info(files)
checkpoint_nums = [get_checkpoint_no(s) for s in files if valid_checkpoint(s)]
if len(checkpoint_nums) > 0:
return os.path.join(args.training_dir, "checkpoint-" + str(max(checkpoint_nums))), max(checkpoint_nums)
return args.init_model_dir, 0
def get_latest_checkpoint(args):
if not os.path.exists(args.training_dir):
return args.init_model_dir, 0
subdirectories = list(next(os.walk(args.training_dir))[1])
def valid_checkpoint(checkpoint):
chk_path = os.path.join(args.training_dir, checkpoint)
scheduler_path = os.path.join(chk_path, "scheduler.pt")
return os.path.exists(scheduler_path)
checkpoint_nums = [get_checkpoint_no(
s) for s in subdirectories if valid_checkpoint(s)]
if len(checkpoint_nums) > 0:
return os.path.join(args.training_dir, "checkpoint-" +
str(max(checkpoint_nums))) + "/", max(checkpoint_nums)
return args.init_model_dir, 0
def get_all_checkpoint(args):
init_list = [args.init_model_dir]
if not os.path.exists(
args.training_dir) or args.inference_one_specified_ckpt:
return init_list
subdirectories = list(next(os.walk(args.training_dir))[1])
def valid_checkpoint(checkpoint):
chk_path = os.path.join(args.training_dir, checkpoint)
scheduler_path = os.path.join(chk_path, "scheduler.pt")
return os.path.exists(scheduler_path)
checkpoint_nums = [
get_checkpoint_no(s) for s in subdirectories if valid_checkpoint(s)
]
checkpoint_nums.sort()
if len(checkpoint_nums) > 0:
return init_list + [
os.path.join(args.training_dir,
"checkpoint-" + str(checkpoint_step))
for checkpoint_step in checkpoint_nums
]
return init_list
def get_latest_checkpoint(args):
if not os.path.exists(args.output_dir):
raise ValueError(
"Output checkpoint directory ({}) not exists. Please remove the --resume_train option"
.format(args.output_dir))
files = list(next(os.walk(args.output_dir))[1])
def valid_checkpoint(checkpoint):
return checkpoint.startswith("checkpoint-")
logger.info("checkpoint files")
logger.info(files)
checkpoint_nums = [
get_checkpoint_no(s) for s in files if valid_checkpoint(s)
]
if len(checkpoint_nums) > 0:
return os.path.join(args.output_dir, "checkpoint-" +
str(max(checkpoint_nums))), max(checkpoint_nums)
else:
raise ValueError(
"Output checkpoint directory ({}) may have some unepected error.".
format(args.output_dir))
def get_latest_checkpoint(args, do_predict_only=False):
if not os.path.exists(args.training_dir):
return args.init_model_dir, 0
files = list(next(os.walk(args.training_dir))[2])
def valid_checkpoint(checkpoint):
return checkpoint.startswith("checkpoint-")
logger.info("checkpoint files")
logger.info(files)
checkpoint_nums = [get_checkpoint_no(s) for s in files if valid_checkpoint(s)]
if do_predict_only:
while (True):
pickle_path = args.output_dir+'passage_'+str(max(checkpoint_nums))+'__emb_p__data_obj_0.pb'
# print(pickle_path, flush=True)
if os.path.exists(pickle_path):
break
else:
checkpoint_nums.remove(max(checkpoint_nums))
if len(checkpoint_nums) > 0:
return os.path.join(args.training_dir, "checkpoint-" + str(max(checkpoint_nums))), max(checkpoint_nums)
return args.init_model_dir, 0
def train(args, model, tokenizer, query_cache, passage_cache):
""" Train the model """
logger.info("Training/evaluation parameters %s", args)
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * \
(torch.distributed.get_world_size() if args.local_rank != -1 else 1)
optimizer_grouped_parameters = []
layer_optim_params = set()
for layer_name in [
"roberta.embeddings", "score_out", "downsample1", "downsample2",
"downsample3"
]:
layer = getattr_recursive(model, layer_name)
if layer is not None:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
if getattr_recursive(model, "roberta.encoder.layer") is not None:
for layer in model.roberta.encoder.layer:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
optimizer_grouped_parameters.append({
"params":
[p for p in model.parameters() if p not in layer_optim_params]
})
if args.optimizer.lower() == "lamb":
optimizer = Lamb(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
elif args.optimizer.lower() == "adamw":
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
else:
raise Exception(
"optimizer {0} not recognized! Can only be lamb or adamW".format(
args.optimizer))
# Check if saved optimizer or scheduler states exist
if os.path.isfile(
os.path.join(args.model_name_or_path,
"optimizer.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train
logger.info("***** Running training *****")
logger.info(" Max steps = %d", args.max_steps)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
global_step = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path):
# set global_step to gobal_step of last saved checkpoint from model
# path
if "-" in args.model_name_or_path:
global_step = int(
args.model_name_or_path.split("-")[-1].split("/")[0])
else:
global_step = 0
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from global step %d", global_step)
tr_loss = 0.0
model.zero_grad()
model.train()
set_seed(args) # Added here for reproductibility
last_ann_no = -1
train_dataloader = None
train_dataloader_iter = None
dev_ndcg = 0
step = 0
save_no = 0
if args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps)
while global_step < args.max_steps:
if step % args.gradient_accumulation_steps == 0 and global_step % args.logging_steps == 0 and global_step % args.save_steps < args.save_steps / 20:
# check if new ann training data is availabe
ann_no, ann_path, ndcg_json = get_latest_ann_data(args.ann_dir)
if ann_path is not None and ann_no != last_ann_no:
logger.info("Training on new add data at %s", ann_path)
with open(ann_path, 'r') as f:
ann_training_data = f.readlines()
dev_ndcg = ndcg_json['ndcg']
ann_checkpoint_path = ndcg_json['checkpoint']
ann_checkpoint_no = get_checkpoint_no(ann_checkpoint_path)
aligned_size = (len(ann_training_data) //
args.world_size) * args.world_size
ann_training_data = ann_training_data[:aligned_size]
logger.info("Total ann queries: %d", len(ann_training_data))
if args.triplet:
train_dataset = StreamingDataset(
ann_training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetTrainingDataProcessingFn(args, query_cache,
passage_cache))
train_dataloader = DataLoader(train_dataset,
batch_size=args.train_batch_size)
train_dataloader_iter = iter(train_dataloader)
# re-warmup
if not args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=len(ann_training_data))
if args.local_rank != -1:
dist.barrier()
if is_first_worker():
# add ndcg at checkpoint step used instead of current step
tb_writer.add_scalar("dev_ndcg", dev_ndcg,
ann_checkpoint_no)
if last_ann_no != -1:
tb_writer.add_scalar("epoch", last_ann_no,
global_step - 1)
tb_writer.add_scalar("epoch", ann_no, global_step)
last_ann_no = ann_no
try:
batch = next(train_dataloader_iter)
except StopIteration:
logger.info("Finished iterating current dataset, begin reiterate")
train_dataloader_iter = iter(train_dataloader)
batch = next(train_dataloader_iter)
batch = tuple(t.to(args.device) for t in batch)
step += 1
if args.triplet:
inputs = {
"query_ids": batch[0].long(),
"attention_mask_q": batch[1].long(),
"input_ids_a": batch[3].long(),
"attention_mask_a": batch[4].long(),
"input_ids_b": batch[6].long(),
"attention_mask_b": batch[7].long()
}
else:
inputs = {
"input_ids_a": batch[0].long(),
"attention_mask_a": batch[1].long(),
"input_ids_b": batch[3].long(),
"attention_mask_b": batch[4].long(),
"labels": batch[6]
}
# sync gradients only at gradient accumulation step
if step % args.gradient_accumulation_steps == 0:
outputs = model(**inputs)
else:
with model.no_sync():
outputs = model(**inputs)
# model outputs are always tuple in transformers (see doc)
loss = outputs[0]
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
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:
if step % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
tr_loss += loss.item()
if step % 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() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
loss_scalar = tr_loss / args.logging_steps
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
logs["loss"] = loss_scalar
tr_loss = 0
if is_first_worker():
for key, value in logs.items():
tb_writer.add_scalar(key, value, global_step)
logger.info(json.dumps({**logs, **{"step": global_step}}))
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(args.output_dir,
"checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
torch.save(args, os.path.join(output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
save_no += 1
if save_no > 1:
ann_no, ann_path, ndcg_json = get_latest_ann_data(
args.ann_dir)
while (ann_no == last_ann_no):
print("Waiting for new ann_data. Sleeping for 1hr!!")
time.sleep(3600)
ann_no, ann_path, ndcg_json = get_latest_ann_data(
args.ann_dir)
dist.barrier()
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step
def train(args, model, tokenizer, query_cache, passage_cache):
""" Train the model """
logger.info("Training/evaluation parameters %s", args)
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(
1, args.n_gpu) #nll loss for query
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * (
torch.distributed.get_world_size() if args.local_rank != -1 else 1)
optimizer = get_optimizer(
args,
model,
weight_decay=args.weight_decay,
)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train!
logger.info("***** Running training *****")
logger.info(" Max steps = %d", args.max_steps)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
tr_loss = 0.0
model.zero_grad()
model.train()
set_seed(args) # Added here for reproductibility
last_ann_no = -1
train_dataloader = None
train_dataloader_iter = None
dev_ndcg = 0
step = 0
iter_count = 0
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps)
global_step = 0
if args.model_name_or_path != "bert-base-uncased":
saved_state = load_states_from_checkpoint(args.model_name_or_path)
global_step = _load_saved_state(model, optimizer, scheduler,
saved_state)
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from global step %d", global_step)
nq_dev_nll_loss, nq_correct_ratio = evaluate_dev(
args, model, passage_cache)
dev_nll_loss_trivia, correct_ratio_trivia = evaluate_dev(
args, model, passage_cache, "-trivia")
if is_first_worker():
tb_writer.add_scalar("dev_nll_loss/dev_nll_loss", nq_dev_nll_loss,
global_step)
tb_writer.add_scalar("dev_nll_loss/correct_ratio",
nq_correct_ratio, global_step)
tb_writer.add_scalar("dev_nll_loss/dev_nll_loss_trivia",
dev_nll_loss_trivia, global_step)
tb_writer.add_scalar("dev_nll_loss/correct_ratio_trivia",
correct_ratio_trivia, global_step)
while global_step < args.max_steps:
if step % args.gradient_accumulation_steps == 0 and global_step % args.logging_steps == 0:
if args.num_epoch == 0:
# check if new ann training data is availabe
ann_no, ann_path, ndcg_json = get_latest_ann_data(args.ann_dir)
if ann_path is not None and ann_no != last_ann_no:
logger.info("Training on new add data at %s", ann_path)
with open(ann_path, 'r') as f:
ann_training_data = f.readlines()
logger.info("Training data line count: %d",
len(ann_training_data))
ann_training_data = [
l for l in ann_training_data
if len(l.split('\t')[2].split(',')) > 1
]
logger.info("Filtered training data line count: %d",
len(ann_training_data))
ann_checkpoint_path = ndcg_json['checkpoint']
ann_checkpoint_no = get_checkpoint_no(ann_checkpoint_path)
aligned_size = (len(ann_training_data) //
args.world_size) * args.world_size
ann_training_data = ann_training_data[:aligned_size]
logger.info("Total ann queries: %d",
len(ann_training_data))
if args.triplet:
train_dataset = StreamingDataset(
ann_training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
train_dataloader = DataLoader(
train_dataset, batch_size=args.train_batch_size)
else:
train_dataset = StreamingDataset(
ann_training_data,
GetTrainingDataProcessingFn(
args, query_cache, passage_cache))
train_dataloader = DataLoader(
train_dataset,
batch_size=args.train_batch_size * 2)
train_dataloader_iter = iter(train_dataloader)
# re-warmup
if not args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=len(ann_training_data))
if args.local_rank != -1:
dist.barrier()
if is_first_worker():
# add ndcg at checkpoint step used instead of current step
tb_writer.add_scalar("retrieval_accuracy/top20_nq",
ndcg_json['top20'],
ann_checkpoint_no)
tb_writer.add_scalar("retrieval_accuracy/top100_nq",
ndcg_json['top100'],
ann_checkpoint_no)
if 'top20_trivia' in ndcg_json:
tb_writer.add_scalar(
"retrieval_accuracy/top20_trivia",
ndcg_json['top20_trivia'], ann_checkpoint_no)
tb_writer.add_scalar(
"retrieval_accuracy/top100_trivia",
ndcg_json['top100_trivia'], ann_checkpoint_no)
if last_ann_no != -1:
tb_writer.add_scalar("epoch", last_ann_no,
global_step - 1)
tb_writer.add_scalar("epoch", ann_no, global_step)
last_ann_no = ann_no
elif step == 0:
train_data_path = os.path.join(args.data_dir, "train-data")
with open(train_data_path, 'r') as f:
training_data = f.readlines()
if args.triplet:
train_dataset = StreamingDataset(
training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
train_dataloader = DataLoader(
train_dataset, batch_size=args.train_batch_size)
else:
train_dataset = StreamingDataset(
training_data,
GetTrainingDataProcessingFn(args, query_cache,
passage_cache))
train_dataloader = DataLoader(
train_dataset, batch_size=args.train_batch_size * 2)
all_batch = [b for b in train_dataloader]
logger.info("Total batch count: %d", len(all_batch))
train_dataloader_iter = iter(train_dataloader)
try:
batch = next(train_dataloader_iter)
except StopIteration:
logger.info("Finished iterating current dataset, begin reiterate")
if args.num_epoch != 0:
iter_count += 1
if is_first_worker():
tb_writer.add_scalar("epoch", iter_count - 1,
global_step - 1)
tb_writer.add_scalar("epoch", iter_count, global_step)
nq_dev_nll_loss, nq_correct_ratio = evaluate_dev(
args, model, passage_cache)
dev_nll_loss_trivia, correct_ratio_trivia = evaluate_dev(
args, model, passage_cache, "-trivia")
if is_first_worker():
tb_writer.add_scalar("dev_nll_loss/dev_nll_loss",
nq_dev_nll_loss, global_step)
tb_writer.add_scalar("dev_nll_loss/correct_ratio",
nq_correct_ratio, global_step)
tb_writer.add_scalar("dev_nll_loss/dev_nll_loss_trivia",
dev_nll_loss_trivia, global_step)
tb_writer.add_scalar("dev_nll_loss/correct_ratio_trivia",
correct_ratio_trivia, global_step)
train_dataloader_iter = iter(train_dataloader)
batch = next(train_dataloader_iter)
dist.barrier()
if args.num_epoch != 0 and iter_count > args.num_epoch:
break
step += 1
if args.triplet:
loss = triplet_fwd_pass(args, model, batch)
else:
loss, correct_cnt = do_biencoder_fwd_pass(args, model, batch)
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
if step % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
tr_loss += loss.item()
if step % 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() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
loss_scalar = tr_loss / args.logging_steps
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
logs["loss"] = loss_scalar
tr_loss = 0
if is_first_worker():
for key, value in logs.items():
tb_writer.add_scalar(key, value, global_step)
logger.info(json.dumps({**logs, **{"step": global_step}}))
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
_save_checkpoint(args, model, optimizer, scheduler,
global_step)
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step
def train(args, model, tokenizer, query_cache, passage_cache):
""" Train the model """
logger.info("Training/evaluation parameters %s", args)
tb_writer = None
if is_first_worker():
tb_writer = SummaryWriter(log_dir=args.log_dir)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
real_batch_size = args.train_batch_size * args.gradient_accumulation_steps * \
(torch.distributed.get_world_size() if args.local_rank != -1 else 1)
optimizer_grouped_parameters = []
layer_optim_params = set()
for layer_name in [
"roberta.embeddings", "score_out", "downsample1", "downsample2",
"downsample3"
]:
layer = getattr_recursive(model, layer_name)
if layer is not None:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
if getattr_recursive(model, "roberta.encoder.layer") is not None:
for layer in model.roberta.encoder.layer:
optimizer_grouped_parameters.append({"params": layer.parameters()})
for p in layer.parameters():
layer_optim_params.add(p)
optimizer_grouped_parameters.append({
"params":
[p for p in model.parameters() if p not in layer_optim_params]
})
if args.optimizer.lower() == "lamb":
optimizer = Lamb(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
elif args.optimizer.lower() == "adamw":
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
else:
raise Exception(
"optimizer {0} not recognized! Can only be lamb or adamW".format(
args.optimizer))
def optimizer_to(optim, device):
for param in optim.state.values():
# Not sure there are any global tensors in the state dict
if isinstance(param, torch.Tensor):
param.data = param.data.to(device)
if param._grad is not None:
param._grad.data = param._grad.data.to(device)
elif isinstance(param, dict):
for subparam in param.values():
if isinstance(subparam, torch.Tensor):
subparam.data = subparam.data.to(device)
if subparam._grad is not None:
subparam._grad.data = subparam._grad.data.to(
device)
torch.cuda.empty_cache()
# Check if saved optimizer or scheduler states exist
if os.path.isfile(
os.path.join(args.model_name_or_path,
"optimizer.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
optimizer.load_state_dict(
torch.load(os.path.join(args.model_name_or_path, "optimizer.pt"),
map_location='cpu'))
optimizer_to(optimizer, args.device)
model.to(args.device)
if args.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True,
)
# Train
logger.info("***** Running training *****")
logger.info(" Max steps = %d", args.max_steps)
logger.info(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size * args.gradient_accumulation_steps *
(torch.distributed.get_world_size() if args.local_rank != -1 else 1),
)
logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
global_step = 0
# Check if continuing training from a checkpoint
if os.path.exists(args.model_name_or_path):
# set global_step to gobal_step of last saved checkpoint from model
# path
if "-" in args.model_name_or_path:
global_step = int(
args.model_name_or_path.split("-")[-1].split("/")[0])
else:
global_step = 0
logger.info(
" Continuing training from checkpoint, will skip to saved global_step"
)
logger.info(" Continuing training from global step %d", global_step)
is_hypersphere_training = (args.hyper_align_weight > 0
or args.hyper_unif_weight > 0)
if is_hypersphere_training:
logger.info(
f"training with hypersphere property regularization, align weight {args.hyper_align_weight}, unif weight {args.hyper_unif_weight}"
)
if not args.dual_training:
args.dual_loss_weight = 0.0
tr_loss_dict = {}
model.zero_grad()
model.train()
set_seed(args) # Added here for reproductibility
last_ann_no = -1
train_dataloader = None
train_dataloader_iter = None
# dev_ndcg = 0
step = 0
if args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=args.max_steps)
if os.path.isfile(os.path.join(
args.model_name_or_path,
"scheduler.pt")) and args.load_optimizer_scheduler:
# Load in optimizer and scheduler states
scheduler.load_state_dict(
torch.load(
os.path.join(args.model_name_or_path, "scheduler.pt")))
while global_step < args.max_steps:
if step % args.gradient_accumulation_steps == 0 and global_step % args.logging_steps == 0:
# check if new ann training data is availabe
ann_no, ann_path, ndcg_json = get_latest_ann_data(
args.ann_dir, is_grouped=(args.grouping_ann_data > 0))
if ann_path is not None and ann_no != last_ann_no:
logger.info("Training on new add data at %s", ann_path)
time.sleep(30) # wait until transmission finished
with open(ann_path, 'r') as f:
ann_training_data = f.readlines()
# marcodev_ndcg = ndcg_json['marcodev_ndcg']
logging.info(f"loading:\n{ndcg_json}")
ann_checkpoint_path = ndcg_json['checkpoint']
ann_checkpoint_no = get_checkpoint_no(ann_checkpoint_path)
aligned_size = (len(ann_training_data) //
args.world_size) * args.world_size
ann_training_data = ann_training_data[:aligned_size]
logger.info(
"Total ann queries: %d",
len(ann_training_data) if args.grouping_ann_data < 0 else
len(ann_training_data) * args.grouping_ann_data)
if args.grouping_ann_data > 0:
if args.polling_loaded_data_batch_from_group:
train_dataset = StreamingDataset(
ann_training_data,
GetGroupedTrainingDataProcessingFn_polling(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetGroupedTrainingDataProcessingFn_origin(
args, query_cache, passage_cache))
else:
if not args.dual_training:
if args.triplet:
train_dataset = StreamingDataset(
ann_training_data,
GetTripletTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
train_dataset = StreamingDataset(
ann_training_data,
GetTrainingDataProcessingFn(
args, query_cache, passage_cache))
else:
# return quadruplet
train_dataset = StreamingDataset(
ann_training_data,
GetQuadrapuletTrainingDataProcessingFn(
args, query_cache, passage_cache))
train_dataloader = DataLoader(train_dataset,
batch_size=args.train_batch_size)
train_dataloader_iter = iter(train_dataloader)
# re-warmup
if not args.single_warmup:
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=args.warmup_steps,
num_training_steps=len(ann_training_data)
if args.grouping_ann_data < 0 else
len(ann_training_data) * args.grouping_ann_data)
if args.local_rank != -1:
dist.barrier()
if is_first_worker():
# add ndcg at checkpoint step used instead of current step
for key in ndcg_json:
if "marcodev" in key:
tb_writer.add_scalar(key, ndcg_json[key],
ann_checkpoint_no)
if 'trec2019_ndcg' in ndcg_json:
tb_writer.add_scalar("trec2019_ndcg",
ndcg_json['trec2019_ndcg'],
ann_checkpoint_no)
if last_ann_no != -1:
tb_writer.add_scalar("epoch", last_ann_no,
global_step - 1)
tb_writer.add_scalar("epoch", ann_no, global_step)
last_ann_no = ann_no
try:
batch = next(train_dataloader_iter)
except StopIteration:
logger.info("Finished iterating current dataset, begin reiterate")
train_dataloader_iter = iter(train_dataloader)
batch = next(train_dataloader_iter)
# original way
if args.grouping_ann_data <= 0:
batch = tuple(t.to(args.device) for t in batch)
if args.triplet:
inputs = {
"query_ids": batch[0].long(),
"attention_mask_q": batch[1].long(),
"input_ids_a": batch[3].long(),
"attention_mask_a": batch[4].long(),
"input_ids_b": batch[6].long(),
"attention_mask_b": batch[7].long()
}
if args.dual_training:
inputs["neg_query_ids"] = batch[9].long()
inputs["attention_mask_neg_query"] = batch[10].long()
inputs["prime_loss_weight"] = args.prime_loss_weight
inputs["dual_loss_weight"] = args.dual_loss_weight
else:
inputs = {
"input_ids_a": batch[0].long(),
"attention_mask_a": batch[1].long(),
"input_ids_b": batch[3].long(),
"attention_mask_b": batch[4].long(),
"labels": batch[6]
}
else:
# the default collate_fn will convert item["q_pos"] into batch format ... I guess
inputs = {
"query_ids": batch["q_pos"][0].to(args.device).long(),
"attention_mask_q": batch["q_pos"][1].to(args.device).long(),
"input_ids_a": batch["d_pos"][0].to(args.device).long(),
"attention_mask_a": batch["d_pos"][1].to(args.device).long(),
"input_ids_b": batch["d_neg"][0].to(args.device).long(),
"attention_mask_b": batch["d_neg"][1].to(args.device).long(),
}
if args.dual_training:
inputs["neg_query_ids"] = batch["q_neg"][0].to(
args.device).long()
inputs["attention_mask_neg_query"] = batch["q_neg"][1].to(
args.device).long()
inputs["prime_loss_weight"] = args.prime_loss_weight
inputs["dual_loss_weight"] = args.dual_loss_weight
inputs["temperature"] = args.temperature
inputs["loss_objective"] = args.loss_objective_function
if is_hypersphere_training:
inputs["alignment_weight"] = args.hyper_align_weight
inputs["uniformity_weight"] = args.hyper_unif_weight
step += 1
if args.local_rank != -1:
# sync gradients only at gradient accumulation step
if step % args.gradient_accumulation_steps == 0:
outputs = model(**inputs)
else:
with model.no_sync():
outputs = model(**inputs)
else:
outputs = model(**inputs)
# model outputs are always tuple in transformers (see doc)
loss = outputs[0]
loss_item_dict = outputs[1]
if args.n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
for k in loss_item_dict:
loss_item_dict[k] = loss_item_dict[k].mean()
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
for k in loss_item_dict:
loss_item_dict[
k] = loss_item_dict[k] / args.gradient_accumulation_steps
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
if args.local_rank != -1:
if step % args.gradient_accumulation_steps == 0:
loss.backward()
else:
with model.no_sync():
loss.backward()
else:
loss.backward()
def incremental_tr_loss(tr_loss_dict, loss_item_dict, total_loss):
for k in loss_item_dict:
if k not in tr_loss_dict:
tr_loss_dict[k] = loss_item_dict[k].item()
else:
tr_loss_dict[k] += loss_item_dict[k].item()
if "loss_total" not in tr_loss_dict:
tr_loss_dict["loss_total"] = total_loss.item()
else:
tr_loss_dict["loss_total"] += total_loss.item()
return tr_loss_dict
tr_loss_dict = incremental_tr_loss(tr_loss_dict,
loss_item_dict,
total_loss=loss)
if step % 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() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.logging_steps > 0 and global_step % args.logging_steps == 0:
logs = {}
learning_rate_scalar = scheduler.get_lr()[0]
logs["learning_rate"] = learning_rate_scalar
for k in tr_loss_dict:
logs[k] = tr_loss_dict[k] / args.logging_steps
tr_loss_dict = {}
if is_first_worker():
for key, value in logs.items():
tb_writer.add_scalar(key, value, global_step)
logger.info(json.dumps({**logs, **{"step": global_step}}))
if is_first_worker(
) and args.save_steps > 0 and global_step % args.save_steps == 0:
# Save model checkpoint
output_dir = os.path.join(args.output_dir,
"checkpoint-{}".format(global_step))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = (
model.module if hasattr(model, "module") else model
) # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
torch.save(args, os.path.join(output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", output_dir)
torch.save(optimizer.state_dict(),
os.path.join(output_dir, "optimizer.pt"))
torch.save(scheduler.state_dict(),
os.path.join(output_dir, "scheduler.pt"))
logger.info("Saving optimizer and scheduler states to %s",
output_dir)
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
tb_writer.close()
return global_step
