def create_and_check_bert_for_pretraining(self, config, input_ids,
token_type_ids, input_mask,
sequence_labels,
token_labels, choice_labels):
model = BertForPreTraining(config=config)
model.to(torch_device)
model.eval()
loss, prediction_scores, seq_relationship_score = model(
input_ids,
attention_mask=input_mask,
token_type_ids=token_type_ids,
masked_lm_labels=token_labels,
next_sentence_label=sequence_labels)
result = {
"loss": loss,
"prediction_scores": prediction_scores,
"seq_relationship_score": seq_relationship_score,
}
self.parent.assertListEqual(
list(result["prediction_scores"].size()),
[self.batch_size, self.seq_length, self.vocab_size])
self.parent.assertListEqual(
list(result["seq_relationship_score"].size()),
[self.batch_size, 2])
self.check_loss_output(result)
def create_and_check_for_pretraining(
self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
):
model = BertForPreTraining(config=config)
model.to(torch_device)
model.eval()
result = model(
input_ids,
attention_mask=input_mask,
token_type_ids=token_type_ids,
labels=token_labels,
next_sentence_label=sequence_labels,
)
self.parent.assertEqual(result.prediction_logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
self.parent.assertEqual(result.seq_relationship_logits.shape, (self.batch_size, 2))
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 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 files for bert pretraining.")
parser.add_argument("--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help="The output directory where the model checkpoints will be written."
)
## Other parameters
# parser.add_argument("--config_name", default="", type=str,
# help="Pretrained config name or path if not the same as model_name")
# parser.add_argument("--tokenizer_name", default="", type=str,
# help="Pretrained tokenizer name or path if not the same as model_name")
parser.add_argument(
"--cache_dir",
default="",
type=str,
help=
"Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument(
"--max_seq_length",
default=128,
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_eval", action="store_true",
# help="Whether to run eval on the dev set.")
# parser.add_argument("--do_predict", action="store_true",
# help="Whether to run predictions on the test set."))
parser.add_argument(
"--do_lower_case",
action="store_true",
help="Set this flag if you are using an uncased model.")
parser.add_argument("--per_gpu_train_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for training.")
parser.add_argument("--per_gpu_eval_batch_size",
default=8,
type=int,
help="Batch size per GPU/CPU for evaluation.")
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("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay",
default=0.0,
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("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--max_steps",
default=-1,
type=int,
help=
"If > 0: set total number of training steps to perform. Override num_train_epochs."
)
parser.add_argument("--warmup_steps",
default=0,
type=int,
help="Linear warmup over warmup_steps.")
parser.add_argument("--logging_steps",
type=int,
default=1000,
help="Log every X updates steps.")
parser.add_argument("--save_steps",
type=int,
default=50,
help="Save checkpoint every X updates steps.")
# parser.add_argument("--eval_all_checkpoints", action="store_true",
# help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
parser.add_argument("--no_cuda",
action="store_true",
help="Avoid using CUDA when available")
parser.add_argument("--overwrite_output_dir",
action="store_true",
help="Overwrite the content of the output directory")
parser.add_argument(
"--overwrite_cache",
action="store_true",
help="Overwrite the cached training and evaluation sets")
parser.add_argument("--seed",
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
"--fp16",
action="store_true",
help=
"Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"
)
parser.add_argument(
"--fp16_opt_level",
type=str,
default="O1",
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("--local_rank",
type=int,
default=-1,
help="For distributed training: local_rank")
parser.add_argument("--server_ip",
type=str,
default="",
help="For distant debugging.")
parser.add_argument("--server_port",
type=str,
default="",
help="For distant debugging.")
parser.add_argument(
"--yago_reference",
action="store_true",
help="Use Reference of Yago types as additional inputs.")
parser.add_argument("--start_task_id", type=int, default=0)
parser.add_argument("--skip_steps", type=int, default=-1)
parser.add_argument("--skip_global_steps", type=int, default=-1)
parser.add_argument("--load_checkpoint", type=str, default="")
args = parser.parse_args()
if 'uncased' in args.model_name_or_path:
args.do_lower_case = True
else:
args.do_lower_case = False
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and not args.overwrite_output_dir:
raise ValueError(
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome."
.format(args.output_dir))
# Setup distant debugging if needed
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
# Setup CUDA, GPU & distributed training
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")
args.n_gpu = torch.cuda.device_count()
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
torch.distributed.init_process_group(backend="nccl")
args.n_gpu = 1
args.device = device
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.warning(
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
args.local_rank, device, args.n_gpu, bool(args.local_rank != -1),
args.fp16)
# Set seed
set_seed(args)
# Use cross entropy ignore index as padding label id so that only real label ids contribute to the loss later
pad_token_label_id = CrossEntropyLoss().ignore_index
# Load pretrained model and tokenizer
if args.local_rank not in [-1, 0]:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
# args.model_type = args.model_type.lower()
config_class, model_class, tokenizer_class = (BertConfig,
BertForPreTraining,
BertTokenizer)
# config = BertConfig() #
bertconfig = BertConfig.from_pretrained(
'bert-base-uncased',
do_lower_case=args.do_lower_case,
cache_dir='/work/smt2/qfeng/Project/huggingface/pretrain/base_uncased')
# tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
# do_lower_case=args.do_lower_case,
# cache_dir=args.cache_dir if args.cache_dir else None)
tokenizer = BertTokenizer.from_pretrained(
'bert-base-uncased',
do_lower_case=True,
cache_dir='/work/smt2/qfeng/Project/huggingface/pretrain/base_uncased')
# model = model_class.from_pretrained(args.model_name_or_path,
# from_tf=bool(".ckpt" in args.model_name_or_path),
# config=config,
# cache_dir=args.cache_dir if args.cache_dir else None)
if args.load_checkpoint == "":
if not args.yago_reference:
config = bertconfig
model = BertForPreTraining(config)
else:
config = YagoRefBertConfig.from_pretrained(
'bert-base-uncased'
if args.do_lower_case else 'bert-base-cased',
reference_size=REFERENCE_SIZE,
cache_dir='/work/smt2/qfeng/Project/huggingface/pretrain/base_{}'
.format('uncased' if args.do_lower_case else 'cased'))
model = YagoRefBertForPreTraining(config)
else:
if 'step' in args.load_checkpoint.split(
'/')[-1] and args.skip_global_steps is not None:
assert (args.load_checkpoint.endswith(str(args.skip_global_steps)))
if not args.yago_reference:
config = BertConfig.from_pretrained(args.load_checkpoint)
model = BertForPreTraining.from_pretrained(args.load_checkpoint)
else:
config = YagoRefBertConfig.from_pretrained(args.load_checkpoint)
model = YagoRefBertForPreTraining.from_pretrained(
args.load_checkpoint)
if args.local_rank == 0:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
model.to(args.device)
logger.info("Training/evaluation parameters %s", args)
# Training
# train_dataset = load_and_cache_examples(args, tokenizer, pad_token_label_id, mode="train") # TODO: need total rewritten
# pickle_list = assign_pickles(args, args.start_task_id)
# train_dataset = load_and_cache_examples(args, tokenizer, pickle_list)
global_step, tr_loss = train(args,
model=model,
tokenizer=tokenizer,
pad_token_label_id=pad_token_label_id)
logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
# Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
if args.do_train and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
# Create output directory if needed
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
os.makedirs(args.output_dir)
logger.info("Saving model checkpoint to %s", args.output_dir)
# Save a trained model, configuration and tokenizer using `save_pretrained()`.
# They can then be reloaded using `from_pretrained()`
model_to_save = model.module if hasattr(
model,
"module") else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)
# Good practice: save your training arguments together with the trained model
torch.save(args, os.path.join(args.output_dir, "training_args.bin"))
# Evaluation TODO: need rewritten:
"""
logger.info("len(sents_list): {}".format(len(sents_list)))
config = BertConfig()
config.num_hidden_layers = 3
config.num_attention_heads = 12
config.hidden_size = 768
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),
from torch.nn import CrossEntropyLoss
from common import AverageMeter
from custom_metrics import LMAccuracy
from data_loader import Data_pretrain
from config import Config
if __name__ == '__main__':
# training_path, file_id, tokenizer, data_name, reduce_memory=False
tokenizer = BertTokenizer.from_pretrained('./bert_pretrain/vocab.txt')
train_data_path = './process_data0.json'
txt = Data_pretrain(train_data_path, tokenizer)
data_iter = DataLoader(txt, shuffle=True, batch_size=2)
bert_config = BertConfig.from_pretrained(Config.config_path)
model = BertForPreTraining(config=bert_config)
model.to(Config.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
}]
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
