418 lines
18 KiB
Python
418 lines
18 KiB
Python
# coding=utf-8
|
|
# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""BERT finetuning runner."""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
|
|
#==================
|
|
import csv
|
|
import os
|
|
import logging
|
|
import argparse
|
|
import random
|
|
import h5py
|
|
from tqdm import tqdm, trange
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, Dataset
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
import math
|
|
from apex import amp
|
|
|
|
|
|
|
|
from tokenization import BertTokenizer
|
|
from modeling import BertForPreTraining, BertConfig
|
|
from optimization import BertAdam, BertAdam_FP16
|
|
|
|
# from fused_adam_local import FusedAdamBert
|
|
from file_utils import PYTORCH_PRETRAINED_BERT_CACHE
|
|
|
|
from apex.optimizers import FusedAdam #, FP16_Optimizer
|
|
#from apex.optimizers import FusedAdam
|
|
from apex.parallel import DistributedDataParallel as DDP
|
|
from schedulers import LinearWarmUpScheduler
|
|
|
|
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s',
|
|
datefmt = '%m/%d/%Y %H:%M:%S',
|
|
level = logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class pretraining_dataset(Dataset):
|
|
|
|
def __init__(self, input_file, max_pred_length):
|
|
self.input_file = input_file
|
|
self.max_pred_length = max_pred_length
|
|
f = h5py.File(input_file, "r")
|
|
self.input_ids = np.asarray(f["input_ids"][:]).astype(np.int64)#[num_instances x max_seq_length])
|
|
self.input_masks = np.asarray(f["input_mask"][:]).astype(np.int64) #[num_instances x max_seq_length]
|
|
self.segment_ids = np.asarray(f["segment_ids"][:]).astype(np.int64) #[num_instances x max_seq_length]
|
|
self.masked_lm_positions = np.asarray(f["masked_lm_positions"][:]).astype(np.int64) #[num_instances x max_pred_length]
|
|
self.masked_lm_ids= np.asarray(f["masked_lm_ids"][:]).astype(np.int64) #[num_instances x max_pred_length]
|
|
self.next_sentence_labels = np.asarray(f["next_sentence_labels"][:]).astype(np.int64) # [num_instances]
|
|
f.close()
|
|
|
|
def __len__(self):
|
|
'Denotes the total number of samples'
|
|
return len(self.input_ids)
|
|
|
|
def __getitem__(self, index):
|
|
|
|
input_ids= torch.from_numpy(self.input_ids[index]) # [max_seq_length]
|
|
input_mask = torch.from_numpy(self.input_masks[index]) #[max_seq_length]
|
|
segment_ids = torch.from_numpy(self.segment_ids[index])# [max_seq_length]
|
|
masked_lm_positions = torch.from_numpy(self.masked_lm_positions[index]) #[max_pred_length]
|
|
masked_lm_ids = torch.from_numpy(self.masked_lm_ids[index]) #[max_pred_length]
|
|
next_sentence_labels = torch.from_numpy(np.asarray(self.next_sentence_labels[index])) #[1]
|
|
|
|
masked_lm_labels = torch.ones(input_ids.shape, dtype=torch.long) * -1
|
|
index = self.max_pred_length
|
|
# store number of masked tokens in index
|
|
if len((masked_lm_positions == 0).nonzero()) != 0:
|
|
index = (masked_lm_positions == 0).nonzero()[0].item()
|
|
masked_lm_labels[masked_lm_positions[:index]] = masked_lm_ids[:index]
|
|
|
|
return [input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels]
|
|
|
|
def main():
|
|
|
|
print("IN NEW MAIN XD\n")
|
|
parser = argparse.ArgumentParser()
|
|
|
|
## Required parameters
|
|
parser.add_argument("--input_dir",
|
|
default=None,
|
|
type=str,
|
|
required=True,
|
|
help="The input data dir. Should contain .hdf5 files for the task.")
|
|
|
|
parser.add_argument("--config_file",
|
|
default=None,
|
|
type=str,
|
|
required=True,
|
|
help="The BERT model config")
|
|
|
|
parser.add_argument("--bert_model", default="bert-large-uncased", type=str,
|
|
help="Bert pre-trained model selected in the list: bert-base-uncased, "
|
|
"bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.")
|
|
|
|
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("--max_seq_length",
|
|
default=512,
|
|
type=int,
|
|
help="The maximum total input sequence length after WordPiece tokenization. \n"
|
|
"Sequences longer than this will be truncated, and sequences shorter \n"
|
|
"than this will be padded.")
|
|
parser.add_argument("--max_predictions_per_seq",
|
|
default=80,
|
|
type=int,
|
|
help="The maximum total of masked tokens in input sequence")
|
|
parser.add_argument("--train_batch_size",
|
|
default=32,
|
|
type=int,
|
|
help="Total batch size for training.")
|
|
parser.add_argument("--learning_rate",
|
|
default=5e-5,
|
|
type=float,
|
|
help="The initial learning rate for Adam.")
|
|
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=1000,
|
|
type=float,
|
|
help="Total number of training steps to perform.")
|
|
parser.add_argument("--warmup_proportion",
|
|
default=0.01,
|
|
type=float,
|
|
help="Proportion of training to perform linear learning rate warmup for. "
|
|
"E.g., 0.1 = 10%% of training.")
|
|
parser.add_argument("--local_rank",
|
|
type=int,
|
|
default=-1,
|
|
help="local_rank for distributed training on gpus")
|
|
parser.add_argument('--seed',
|
|
type=int,
|
|
default=42,
|
|
help="random seed for initialization")
|
|
parser.add_argument('--gradient_accumulation_steps',
|
|
type=int,
|
|
default=1,
|
|
help="Number of updates steps to accumualte before performing a backward/update pass.")
|
|
parser.add_argument('--fp16',
|
|
default=False,
|
|
action='store_true',
|
|
help="Whether to use 16-bit float precision instead of 32-bit")
|
|
parser.add_argument('--loss_scale',
|
|
type=float, default=0.0,
|
|
help='Loss scaling, positive power of 2 values can improve fp16 convergence.')
|
|
parser.add_argument('--log_freq',
|
|
type=float, default=10.0,
|
|
help='frequency of logging loss.')
|
|
parser.add_argument('--checkpoint_activations',
|
|
default=False,
|
|
action='store_true',
|
|
help="Whether to use gradient checkpointing")
|
|
parser.add_argument("--resume_from_checkpoint",
|
|
default=False,
|
|
action='store_true',
|
|
help="Whether to resume training from checkpoint.")
|
|
parser.add_argument('--resume_step',
|
|
type=int,
|
|
default=-1,
|
|
help="Step to resume training from.")
|
|
parser.add_argument('--num_steps_per_checkpoint',
|
|
type=int,
|
|
default=2000,
|
|
help="Number of update steps until a model checkpoint is saved to disk.")
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
|
|
assert(torch.cuda.is_available())
|
|
|
|
if args.local_rank == -1:
|
|
device = torch.device("cuda")
|
|
n_gpu = torch.cuda.device_count()
|
|
else:
|
|
torch.cuda.set_device(args.local_rank)
|
|
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', init_method='env://')
|
|
|
|
logger.info("device %s n_gpu %d distributed training %r", device, n_gpu, bool(args.local_rank != -1))
|
|
|
|
if args.gradient_accumulation_steps < 1:
|
|
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
|
|
args.gradient_accumulation_steps))
|
|
if args.train_batch_size % args.gradient_accumulation_steps != 0:
|
|
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, batch size {} should be divisible".format(
|
|
args.gradient_accumulation_steps, args.train_batch_size))
|
|
|
|
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
|
|
|
|
|
|
|
|
if not args.resume_from_checkpoint and os.path.exists(args.output_dir) and (os.listdir(args.output_dir) and os.listdir(args.output_dir)!=['logfile.txt']):
|
|
raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
|
|
|
|
if not args.resume_from_checkpoint:
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
# Prepare model
|
|
config = BertConfig.from_json_file(args.config_file)
|
|
model = BertForPreTraining(config)
|
|
|
|
|
|
if not args.resume_from_checkpoint:
|
|
global_step = 0
|
|
else:
|
|
if args.resume_step == -1:
|
|
model_names = [f for f in os.listdir(args.output_dir) if f.endswith(".pt")]
|
|
args.resume_step = max([int(x.split('.pt')[0].split('_')[1].strip()) for x in model_names])
|
|
|
|
global_step = args.resume_step
|
|
|
|
checkpoint = torch.load(os.path.join(args.output_dir, "ckpt_{}.pt".format(global_step)), map_location="cpu")
|
|
model.load_state_dict(checkpoint['model'], strict=False)
|
|
|
|
print("resume step from ", args.resume_step)
|
|
|
|
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}
|
|
]
|
|
|
|
|
|
|
|
if args.fp16:
|
|
|
|
optimizer = FusedAdam(optimizer_grouped_parameters,
|
|
lr=args.learning_rate,
|
|
#warmup=args.warmup_proportion,
|
|
#t_total=args.max_steps,
|
|
bias_correction=False,
|
|
weight_decay=0.01,
|
|
max_grad_norm=1.0)
|
|
|
|
if args.loss_scale == 0:
|
|
# optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
|
|
model, optimizer = amp.initialize(model, optimizer, opt_level="O2", keep_batchnorm_fp32=False, loss_scale="dynamic")
|
|
else:
|
|
# optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale)
|
|
model, optimizer = amp.initialize(model, optimizer, opt_level="O2", keep_batchnorm_fp32=False, loss_scale=args.loss_scale)
|
|
|
|
scheduler = LinearWarmUpScheduler(optimizer, warmup=args.warmup_proportion, total_steps=args.max_steps)
|
|
|
|
else:
|
|
optimizer = BertAdam(optimizer_grouped_parameters,
|
|
lr=args.learning_rate,
|
|
warmup=args.warmup_proportion,
|
|
t_total=args.max_steps)
|
|
|
|
|
|
|
|
if args.resume_from_checkpoint:
|
|
optimizer.load_state_dict(checkpoint['optimizer']) # , strict=False)
|
|
|
|
|
|
|
|
if args.local_rank != -1:
|
|
model = DDP(model)
|
|
elif n_gpu > 1:
|
|
model = torch.nn.DataParallel(model)
|
|
|
|
|
|
files = [os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if os.path.isfile(os.path.join(args.input_dir, f))]
|
|
files.sort()
|
|
|
|
num_files = len(files)
|
|
|
|
|
|
logger.info("***** Running training *****")
|
|
# logger.info(" Num examples = %d", len(train_data))
|
|
logger.info(" Batch size = %d", args.train_batch_size)
|
|
print(" LR = ", args.learning_rate)
|
|
|
|
|
|
model.train()
|
|
print("Training. . .")
|
|
|
|
most_recent_ckpts_paths = []
|
|
|
|
print("Training. . .")
|
|
tr_loss = 0.0 # total added training loss
|
|
average_loss = 0.0 # averaged loss every args.log_freq steps
|
|
epoch = 0
|
|
training_steps = 0
|
|
while True:
|
|
if not args.resume_from_checkpoint:
|
|
random.shuffle(files)
|
|
f_start_id = 0
|
|
else:
|
|
f_start_id = checkpoint['files'][0]
|
|
files = checkpoint['files'][1:]
|
|
args.resume_from_checkpoint = False
|
|
for f_id in range(f_start_id, len(files)):
|
|
data_file = files[f_id]
|
|
logger.info("file no %s file %s" %(f_id, data_file))
|
|
train_data = pretraining_dataset(input_file=data_file, max_pred_length=args.max_predictions_per_seq)
|
|
|
|
if args.local_rank == -1:
|
|
train_sampler = RandomSampler(train_data)
|
|
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size * n_gpu, num_workers=4, pin_memory=True)
|
|
else:
|
|
train_sampler = DistributedSampler(train_data)
|
|
|
|
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size, num_workers=4, pin_memory=True)
|
|
for step, batch in enumerate(tqdm(train_dataloader, desc="File Iteration")):
|
|
|
|
training_steps += 1
|
|
batch = [t.to(device) for t in batch]
|
|
input_ids, segment_ids, input_mask, masked_lm_labels, next_sentence_labels = batch#\
|
|
loss = model(input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask, masked_lm_labels=masked_lm_labels, next_sentence_label=next_sentence_labels, checkpoint_activations=args.checkpoint_activations)
|
|
if 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:
|
|
# optimizer.backward(loss)
|
|
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
|
scaled_loss.backward()
|
|
else:
|
|
loss.backward()
|
|
tr_loss += loss
|
|
average_loss += loss.item()
|
|
|
|
if training_steps % args.gradient_accumulation_steps == 0:
|
|
if args.fp16:
|
|
scheduler.step()
|
|
optimizer.step()
|
|
optimizer.zero_grad()
|
|
global_step += 1
|
|
|
|
|
|
|
|
if training_steps == 1 * args.gradient_accumulation_steps:
|
|
logger.info("Step:{} Average Loss = {} Step Loss = {} LR {}".format(global_step, average_loss,
|
|
loss.item(), optimizer.param_groups[0]['lr']))
|
|
|
|
if training_steps % (args.log_freq * args.gradient_accumulation_steps) == 0:
|
|
logger.info("Step:{} Average Loss = {} Step Loss = {} LR {}".format(global_step, average_loss / args.log_freq,
|
|
loss.item(), optimizer.param_groups[0]['lr']))
|
|
average_loss = 0
|
|
|
|
|
|
if global_step >= args.max_steps or training_steps == 1 * args.gradient_accumulation_steps or training_steps % (args.num_steps_per_checkpoint * args.gradient_accumulation_steps) == 0:
|
|
if (not torch.distributed.is_initialized() or (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0)):
|
|
# Save a trained model
|
|
logger.info("** ** * Saving fine - tuned model ** ** * ")
|
|
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self
|
|
output_save_file = os.path.join(args.output_dir, "ckpt_{}.pt".format(global_step))
|
|
|
|
torch.save({'model' : model_to_save.state_dict(),
|
|
'optimizer' : optimizer.state_dict(),
|
|
'files' : [f_id] + files }, output_save_file)
|
|
|
|
most_recent_ckpts_paths.append(output_save_file)
|
|
if len(most_recent_ckpts_paths) > 3:
|
|
ckpt_to_be_removed = most_recent_ckpts_paths.pop(0)
|
|
os.remove(ckpt_to_be_removed)
|
|
|
|
if global_step >= args.max_steps:
|
|
tr_loss = tr_loss * args.gradient_accumulation_steps / training_steps
|
|
if (torch.distributed.is_initialized()):
|
|
tr_loss /= torch.distributed.get_world_size()
|
|
torch.distributed.all_reduce(tr_loss)
|
|
logger.info("Total Steps:{} Final Loss = {}".format(training_steps, tr_loss.item()))
|
|
return
|
|
del train_dataloader
|
|
del train_sampler
|
|
del train_data
|
|
#for obj in gc.get_objects():
|
|
# if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
|
|
# del obj
|
|
|
|
torch.cuda.empty_cache()
|
|
epoch += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|