config.intermediate_size = 3072
config.max_position_embeddings = 512
config.vocab_size = 32000
logger.info("USE_NSP: {}".format(USE_NSP))
if USE_NSP:
model = BertForPreTraining(config)
else:
model = BertForPreTrainingWithoutNSP(config)
model.to(device)
logger.info(config)
logger.info(model)
optimizer = AdamW(model.parameters(), lr=2e-5)
model.train()
train_losses = []
for i in range(1, MAX_STEPS + 1):
optimizer.zero_grad()
sent_pairs = create_sent_pairs(sents_list, batch_size=BATCH_SIZE)
encoded = encode_sent_pairs(sent_pairs)
res = model(
encoded["input_ids"].to(device),
token_type_ids=None,
attention_mask=encoded["attention_mask"].to(device),
labels=encoded["labels"].to(device),
next_sentence_label=encoded["next_sentence_label"].to(device),
)
loss = res.loss
if i % 100 == 0:
logger.info("training step {}, loss {}".format(i, loss))
def main():
parser = ArgumentParser()
parser.add_argument('--data_name', default='albert', type=str)
parser.add_argument(
"--file_num",
type=int,
default=10,
help="Number of dynamic masking to pregenerate (with different masks)"
) # TODO: --file-num 이름이 어울리지 않음
parser.add_argument(
"--reduce_memory",
action="store_true",
help=
"Store training data as on-disc memmaps to massively reduce memory usage"
)
parser.add_argument("--epochs",
type=int,
default=4,
help="Number of epochs to train for")
parser.add_argument('--share_type',
default='all',
type=str,
choices=['all', 'attention', 'ffn', 'None'])
parser.add_argument('--num_eval_steps', default=100)
parser.add_argument('--num_save_steps', default=200)
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
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("--train_batch_size",
default=4,
type=int,
help="Total batch size for training.")
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("--warmup_proportion",
default=0.1,
type=float,
help="Linear warmup over warmup_steps.")
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)
parser.add_argument("--learning_rate",
default=0.00176,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--fp16_opt_level',
type=str,
default='O2',
help=
"For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
"See details at https://nvidia.github.io/apex/amp.html")
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
args = parser.parse_args()
pregenerated_data = config['data_dir'] / "corpus/train"
assert pregenerated_data.is_dir(), \
"--pregenerated_data should point to the folder of files made by prepare_lm_data_mask.py!"
samples_per_epoch = 0
for i in range(args.file_num):
data_file = pregenerated_data / f"{args.data_name}_file_{i}.json"
metrics_file = pregenerated_data / f"{args.data_name}_file_{i}_metrics.json" # TODO: 어디에 있지?
if data_file.is_file() and metrics_file.is_file():
metrics = json.loads(metrics_file.read_text())
samples_per_epoch += metrics['num_training_examples']
else:
if i == 0:
exit("No training data was found!")
print(
f"Warning! There are fewer epochs of pregenerated data ({i}) than training epochs ({args.epochs})."
)
print(
"This script will loop over the available data, but training diversity may be negatively impacted."
)
break
logger.info(f"samples_per_epoch: {samples_per_epoch}")
if args.local_rank == -1 or args.no_cuda:
device = torch.device(f"cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
args.n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
f"device: {device} , distributed training: {bool(args.local_rank != -1)}, 16-bits training: {args.fp16}, "
f"share_type: {args.share_type}")
if args.gradient_accumulation_steps < 1:
raise ValueError(
f"Invalid gradient_accumulation_steps parameter: {args.gradient_accumulation_steps}, should be >= 1"
)
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
seed_everything(args.seed)
tokenizer = BertTokenizer(vocab_file=config['albert_vocab_path'])
total_train_examples = samples_per_epoch * args.epochs
num_train_optimization_steps = int(total_train_examples /
args.train_batch_size /
args.gradient_accumulation_steps)
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
args.warmup_steps = int(num_train_optimization_steps *
args.warmup_proportion)
bert_config = BertConfig.from_pretrained(str(config['albert_config_path']),
share_type=args.share_type)
model = BertForPreTraining(config=bert_config)
# model = BertForMaskedLM.from_pretrained(config['checkpoint_dir'] / 'checkpoint-580000')
model.to(device)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.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
}]
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
# optimizer = Lamb(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
lr_scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=num_train_optimization_steps)
if args.fp16:
try:
from apex import amp
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args.fp16_opt_level)
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
if args.n_gpu > 1:
model = torch.nn.DataParallel(model)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
global_step = 0
mask_metric = LMAccuracy()
sop_metric = LMAccuracy()
tr_mask_acc = AverageMeter()
tr_sop_acc = AverageMeter()
tr_loss = AverageMeter()
tr_mask_loss = AverageMeter()
tr_sop_loss = AverageMeter()
loss_fct = CrossEntropyLoss(ignore_index=-1)
train_logs = {}
logger.info("***** Running training *****")
logger.info(f" Num examples = {total_train_examples}")
logger.info(f" Batch size = {args.train_batch_size}")
logger.info(f" Num steps = {num_train_optimization_steps}")
logger.info(f" warmup_steps = {args.warmup_steps}")
start_time = time.time()
seed_everything(args.seed) # Added here for reproducibility
for epoch in range(args.epochs):
for idx in range(args.file_num):
epoch_dataset = PregeneratedDataset(
file_id=idx,
training_path=pregenerated_data,
tokenizer=tokenizer,
reduce_memory=args.reduce_memory,
data_name=args.data_name)
if args.local_rank == -1:
train_sampler = RandomSampler(epoch_dataset)
else:
train_sampler = DistributedSampler(epoch_dataset)
train_dataloader = DataLoader(epoch_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, lm_label_ids, is_next = batch
outputs = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
prediction_scores = outputs[0]
seq_relationship_score = outputs[1]
masked_lm_loss = loss_fct(
prediction_scores.view(-1, bert_config.vocab_size),
lm_label_ids.view(-1))
next_sentence_loss = loss_fct(
seq_relationship_score.view(-1, 2), is_next.view(-1))
loss = masked_lm_loss + next_sentence_loss
mask_metric(logits=prediction_scores.view(
-1, bert_config.vocab_size),
target=lm_label_ids.view(-1))
sop_metric(logits=seq_relationship_score.view(-1, 2),
target=is_next.view(-1))
if args.n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
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()
nb_tr_steps += 1
tr_mask_acc.update(mask_metric.value(), n=input_ids.size(0))
tr_sop_acc.update(sop_metric.value(), n=input_ids.size(0))
tr_loss.update(loss.item(), n=1)
tr_mask_loss.update(masked_lm_loss.item(), n=1)
tr_sop_loss.update(next_sentence_loss.item(), n=1)
if (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)
lr_scheduler.step()
optimizer.step()
optimizer.zero_grad()
global_step += 1
if global_step % args.num_eval_steps == 0:
now = time.time()
eta = now - start_time
if eta > 3600:
eta_format = ('%d:%02d:%02d' %
(eta // 3600,
(eta % 3600) // 60, eta % 60))
elif eta > 60:
eta_format = '%d:%02d' % (eta // 60, eta % 60)
else:
eta_format = '%ds' % eta
train_logs['loss'] = tr_loss.avg
train_logs['mask_acc'] = tr_mask_acc.avg
train_logs['sop_acc'] = tr_sop_acc.avg
train_logs['mask_loss'] = tr_mask_loss.avg
train_logs['sop_loss'] = tr_sop_loss.avg
show_info = f'[Training]:[{epoch}/{args.epochs}]{global_step}/{num_train_optimization_steps} ' \
f'- ETA: {eta_format}' + "-".join(
[f' {key}: {value:.4f} ' for key, value in train_logs.items()])
logger.info(show_info)
tr_mask_acc.reset()
tr_sop_acc.reset()
tr_loss.reset()
tr_mask_loss.reset()
tr_sop_loss.reset()
start_time = now
if global_step % args.num_save_steps == 0:
if args.local_rank in [-1, 0] and args.num_save_steps > 0:
# Save model checkpoint
output_dir = config[
'checkpoint_dir'] / f'lm-checkpoint-{global_step}'
if not output_dir.exists():
output_dir.mkdir()
# save model
model_to_save = model.module if hasattr(
model, 'module'
) else model # Take care of distributed/parallel training
model_to_save.save_pretrained(str(output_dir))
torch.save(args, str(output_dir / 'training_args.bin'))
logger.info("Saving model checkpoint to %s",
output_dir)
# save config
output_config_file = output_dir / CONFIG_NAME
with open(str(output_config_file), 'w') as f:
f.write(model_to_save.config.to_json_string())
# save vocab
tokenizer.save_vocabulary(output_dir)
def train(args):
if not os.path.exists(args.save_dir): os.mkdir(args.save_dir)
if args.gpu != '-1' and torch.cuda.is_available():
device = torch.device('cuda')
torch.cuda.set_rng_state(torch.cuda.get_rng_state())
torch.backends.cudnn.deterministic = True
else:
device = torch.device('cpu')
config = {
'train': {
'unchanged_variable_weight': 0.1,
'buffer_size': 5000
},
'encoder': {
'type': 'SequentialEncoder'
},
'data': {
'vocab_file': 'data/vocab.bpe10000/vocab'
}
}
train_set = Dataset('data/preprocessed_data/train-shard-*.tar')
dev_set = Dataset('data/preprocessed_data/dev.tar')
vocab = Vocab.load('data/vocab.bpe10000/vocab')
if args.decoder:
vocab_size = len(vocab.all_subtokens) + 1
else:
vocab_size = len(vocab.source_tokens) + 1
max_iters = args.max_iters
lr = args.lr
warm_up = args.warm_up
batch_size = 4096
effective_batch_size = args.batch_size
max_embeds = 1000 if args.decoder else 512
bert_config = BertConfig(vocab_size=vocab_size,
max_position_embeddings=max_embeds,
num_hidden_layers=6,
hidden_size=256,
num_attention_heads=4)
model = BertForPreTraining(bert_config)
if args.restore:
state_dict = torch.load(os.path.join(args.save_dir, args.res_name))
model.load_state_dict(state_dict['model'])
batch_count = state_dict['step']
epoch = state_dict['epoch']
model.train()
model.to(device)
if len(args.gpu) > 1 and device == torch.device('cuda'):
model = nn.DataParallel(model)
def lr_func(step):
if step > warm_up:
return (max_iters - step) / (max_iters - warm_up)
else:
return (step / warm_up)
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
eps=1e-6,
weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer,
lr_lambda=lr_func,
last_epoch=-1)
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction='none')
if args.restore:
optimizer.load_state_dict(state_dict['optim'])
scheduler.load_state_dict(state_dict['scheduler'])
batch_count = 0
epoch = 0
cum_loss = 0.0
while True:
# load training dataset, which is a collection of ASTs and maps of gold-standard renamings
train_set_iter = train_set.batch_iterator(
batch_size=batch_size,
return_examples=False,
config=config,
progress=True,
train=True,
max_seq_len=512,
num_readers=args.num_readers,
num_batchers=args.num_batchers)
epoch += 1
print("Epoch {}".format(epoch))
loss = 0
num_seq = 0
optimizer.zero_grad()
for batch in train_set_iter:
if args.decoder:
input_ids = batch.tensor_dict['prediction_target'][
'src_with_true_var_names']
else:
input_ids = batch.tensor_dict['src_code_tokens']
attention_mask = torch.ones_like(input_ids)
attention_mask[input_ids == 0] = 0.0
assert torch.max(input_ids) < vocab_size
assert torch.min(input_ids) >= 0
if input_ids.shape[0] > max_embeds:
print(
"Warning - length {} is greater than max length {}. Skipping."
.format(input_ids.shape[0], max_embeds))
continue
input_ids, labels = mask_tokens(inputs=input_ids,
mask_token_id=vocab_size - 1,
vocab_size=vocab_size,
mlm_probability=0.15)
input_ids[attention_mask == 0] = 0
labels[attention_mask == 0] = -100
if torch.cuda.is_available():
input_ids = input_ids.cuda()
labels = labels.cuda()
attention_mask = attention_mask.cuda()
outputs = model(input_ids=input_ids,
attention_mask=attention_mask,
masked_lm_labels=labels)
unreduced_loss = loss_fn(
outputs[0].view(-1, bert_config.vocab_size),
labels.view(-1)).reshape(labels.shape) / (
torch.sum(labels != -100, axis=1).unsqueeze(1) + 1e-7)
loss += unreduced_loss.sum()
num_seq += input_ids.shape[0]
if num_seq > effective_batch_size:
batch_count += 1
loss /= num_seq
cum_loss += loss.item()
if batch_count % 20 == 0:
print("{} batches, Loss : {:.4}, LR : {:.6}".format(
batch_count, cum_loss / 20,
scheduler.get_lr()[0]))
cum_loss = 0.0
if batch_count % 10000 == 0:
fname1 = os.path.join(
args.save_dir, 'bert_{}_step_{}.pth'.format(
('decoder' if args.decoder else 'encoder'),
batch_count))
fname2 = os.path.join(
args.save_dir, 'bert_{}.pth'.format(
('decoder' if args.decoder else 'encoder'),
batch_count))
state = {
'epoch': epoch,
'step': batch_count,
'model': model.module.state_dict(),
'optim': optimizer.state_dict(),
'scheduler': scheduler.state_dict()
}
torch.save(state, fname1)
torch.save(state, fname2)
print("Saved file to path {}".format(fname1))
print("Saved file to path {}".format(fname2))
loss.backward()
optimizer.step()
scheduler.step()
optimizer.zero_grad()
loss = 0
num_seq = 0
if batch_count == max_iters:
print(f'[Learner] Reached max iters', file=sys.stderr)
exit()
print("Max_len = {}".format(max_len))
break
