def train(args, train_dataset, model, tokenizer):
""" Train the model """
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter(log_dir=args.output_dir+'/runs/tensorboard')
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight', 'head_weights', 'layer_weights']
optimizer_grouped_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}
]
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
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(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
set_seed(args) # Added here for reproductibility (even between python 2 and 3)
best_methods = {}
for _ in train_iterator:
epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
for step, batch in enumerate(epoch_iterator):
model.train()
batch = tuple(t.to(args.device) for t in batch)
inputs = {'input_ids': batch[0],
'attention_mask': batch[1],
'token_type_ids': batch[2] if args.model_type in ['bert', 'xlnet'] else None, # XLM don't use segment_ids
'labels': batch[3]}
outputs = model(**inputs)
loss = outputs[0] # model outputs are always tuple in pytorch-transformers (see doc)
if args.n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu parallel training
if args.smart_head and args.regularize_head:
tmp_model = model.module if hasattr(model, 'module') else model
loss = loss + smart_mask_regulaziation(tmp_model.bert.head_weights, args.shm_reg_type,
args.shm_reg, args.device, tmp_model.bert.eps,
tmp_model.bert.temp, tmp_model.bert.gamma_zeta_ratio)
if args.smart_pooling and args.regularize_pooling:
tmp_model = model.module if hasattr(model, 'module') else model
loss = loss + smart_mask_regulaziation(tmp_model.bert.layer_weights, args.lp_reg_type,
args.lp_reg, args.device, tmp_model.bert.eps,
tmp_model.bert.temp, tmp_model.bert.gamma_zeta_ratio)
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()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
scheduler.step() # Update learning rate schedule
optimizer.step()
model.zero_grad()
global_step += 1
if args.smart_head:
tmp_model = model.module if hasattr(model, 'module') else model
mask = tmp_model.bert.gate_mask(tmp_model.bert.head_weights, args.shm_reg_type,
tmp_model.bert.head_mask_size, args.device)
exact_sparsity, relaxed_sparsity = sparsity_rate(mask, args.device)
tb_writer.add_scalar('exact_sparsity', exact_sparsity, global_step)
tb_writer.add_scalar('relaxed_sparsity', relaxed_sparsity, global_step)
tb_writer.add_histogram('head_mask_weight', tmp_model.bert.head_weights, global_step)
if args.smart_pooling:
tmp_model = model.module if hasattr(model, 'module') else model
mask = tmp_model.bert.gate_mask(tmp_model.bert.layer_weights, args.lp_reg_type,
tmp_model.bert.layer_mask_size, args.device)
for i in range(tmp_model.bert.config.num_hidden_layers):
tb_writer.add_scalar('layer_weight_%d'%i, mask[i], global_step)
if args.local_rank in [-1, 0] and ((args.logging_steps > 0 and global_step % args.logging_steps == 0) or (args.save_steps > 0 and global_step % args.save_steps == 0)):
# Log metrics
if args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer)
for key, value in results.items():
tb_writer.add_scalar('eval_{}'.format(key), value, global_step)
if key in best_methods:
if value > best_methods[key]:
output_dir = os.path.join(args.output_dir, 'checkpoint-best-{}'.format(str(key)))
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)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
best_methods[key] = value
else:
output_dir = os.path.join(args.output_dir, 'checkpoint-best-{}'.format(str(key)))
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)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
best_methods[key] = value
tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
tb_writer.add_scalar('loss', (tr_loss - logging_loss)/args.logging_steps, global_step)
logging_loss = tr_loss
if args.max_steps > 0 and global_step > args.max_steps:
epoch_iterator.close()
break
if args.local_rank in [-1, 0]: # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer)
for key, value in results.items():
if key in best_methods:
if value > best_methods[key]:
output_dir = os.path.join(args.output_dir, 'checkpoint-best-{}'.format(str(key)))
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)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
best_methods[key] = value
else:
output_dir = os.path.join(args.output_dir, 'checkpoint-best-{}'.format(str(key)))
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)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
if args.max_steps > 0 and global_step > args.max_steps:
train_iterator.close()
break
if args.local_rank in [-1, 0]:
tb_writer.close()
return global_step, tr_loss / global_step
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model_name',
type=str,
default='openai-gpt',
help='pretrained model name')
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument('--train_dataset', type=str, default='')
parser.add_argument('--eval_dataset', type=str, default='')
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--num_train_epochs', type=int, default=3)
parser.add_argument('--train_batch_size', type=int, default=8)
parser.add_argument('--eval_batch_size', type=int, default=16)
parser.add_argument("--adam_epsilon",
default=1e-8,
type=float,
help="Epsilon for Adam optimizer.")
parser.add_argument('--max_grad_norm', type=int, default=1)
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('--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', type=float, default=6.25e-5)
parser.add_argument("--warmup_steps",
default=0,
type=int,
help="Linear warmup over warmup_steps.")
parser.add_argument('--lr_schedule', type=str, default='warmup_linear')
parser.add_argument('--weight_decay', type=float, default=0.01)
parser.add_argument('--lm_coef', type=float, default=0.9)
parser.add_argument('--n_valid', type=int, default=374)
parser.add_argument('--server_ip',
type=str,
default='',
help="Can be used for distant debugging.")
parser.add_argument('--server_port',
type=str,
default='',
help="Can be used for distant debugging.")
args = parser.parse_args()
print(args)
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()
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = torch.cuda.device_count()
logger.info("device: {}, n_gpu {}".format(device, n_gpu))
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# Load tokenizer and model
# This loading functions also add new tokens and embeddings called `special tokens`
# These new embeddings will be fine-tuned on the RocStories dataset
special_tokens = ['_start_', '_delimiter_', '_classify_']
tokenizer = OpenAIGPTTokenizer.from_pretrained(args.model_name)
tokenizer.add_tokens(new_tokens=special_tokens)
special_tokens_ids = tokenizer.convert_tokens_to_ids(tokens=special_tokens)
model = OpenAIGPTDoubleHeadsModel.from_pretrained(
args.model_name, n_special=len(special_tokens))
model.to(device)
# Load and encode the datasets
if not args.train_dataset and not args.eval_dataset:
roc_stories = cached_path(ROCSTORIES_URL)
def tokenize_and_encode(obj):
""" Tokenize and encode a nested object """
if isinstance(obj, str):
return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(obj))
elif isinstance(obj, int):
return obj
return list(tokenize_and_encode(o) for o in obj)
logger.info("Encoding dataset...")
train_dataset = load_rocstories_dataset(args.train_dataset)
eval_dataset = load_rocstories_dataset(args.eval_dataset)
datasets = (train_dataset, eval_dataset)
encoded_datasets = tokenize_and_encode(datasets)
# Compute the max input length for the Transformer
max_length = model.config.n_positions // 2 - 2
input_length = max(len(story[:max_length]) + max(len(cont1[:max_length]), len(cont2[:max_length])) + 3 \
for dataset in encoded_datasets for story, cont1, cont2, _ in dataset)
input_length = min(input_length, model.config.n_positions
) # Max size of input for the pre-trained model
# Prepare inputs tensors and dataloaders
tensor_datasets = pre_process_datasets(encoded_datasets, input_length,
max_length, *special_tokens_ids)
train_tensor_dataset, eval_tensor_dataset = tensor_datasets[
0], tensor_datasets[1]
train_data = TensorDataset(*train_tensor_dataset)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
eval_data = TensorDataset(*eval_tensor_dataset)
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
# Prepare optimizer
if args.do_train:
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps //\
(len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(train_dataloader)\
// args.gradient_accumulation_steps * args.num_train_epochs
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':
args.weight_decay
}, {
'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)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=t_total)
if args.do_train:
nb_tr_steps, tr_loss, exp_average_loss = 0, 0, None
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_steps = 0
tqdm_bar = tqdm(train_dataloader, desc="Training")
for step, batch in enumerate(tqdm_bar):
batch = tuple(t.to(device) for t in batch)
input_ids, mc_token_ids, lm_labels, mc_labels = batch
losses = model(input_ids, mc_token_ids, lm_labels, mc_labels)
loss = args.lm_coef * losses[0] + losses[1]
loss.backward()
scheduler.step()
optimizer.step()
optimizer.zero_grad()
tr_loss += loss.item()
exp_average_loss = loss.item(
) if exp_average_loss is None else 0.7 * exp_average_loss + 0.3 * loss.item(
)
nb_tr_steps += 1
tqdm_bar.desc = "Training loss: {:.2e} lr: {:.2e}".format(
exp_average_loss,
scheduler.get_lr()[0])
# Save a trained model
if args.do_train:
# Save a trained model, configuration and tokenizer
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(args.output_dir)
# Load a trained model and vocabulary that you have fine-tuned
model = OpenAIGPTDoubleHeadsModel.from_pretrained(args.output_dir)
tokenizer = OpenAIGPTTokenizer.from_pretrained(args.output_dir)
model.to(device)
if args.do_eval:
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for batch in tqdm(eval_dataloader, desc="Evaluating"):
batch = tuple(t.to(device) for t in batch)
input_ids, mc_token_ids, lm_labels, mc_labels = batch
with torch.no_grad():
_, mc_loss, _, mc_logits = model(input_ids, mc_token_ids,
lm_labels, mc_labels)
mc_logits = mc_logits.detach().cpu().numpy()
mc_labels = mc_labels.to('cpu').numpy()
tmp_eval_accuracy = accuracy(mc_logits, mc_labels)
eval_loss += mc_loss.mean().item()
eval_accuracy += tmp_eval_accuracy
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = eval_accuracy / nb_eval_examples
train_loss = tr_loss / nb_tr_steps if args.do_train else None
result = {
'eval_loss': eval_loss,
'eval_accuracy': eval_accuracy,
'train_loss': train_loss
}
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
def train(self,
model,
model_name,
B, N_for_train, N_for_eval, K, Q,
na_rate=0,
learning_rate=1e-1,
lr_step_size=20000,
weight_decay=1e-5,
train_iter=30000,
val_iter=1000,
val_step=2000,
test_iter=3000,
load_ckpt=None,
save_ckpt=None,
pytorch_optim=optim.SGD,
bert_optim=False,
warmup=True,
warmup_step=300,
grad_iter=1,
fp16=False,
pair=False,
adv_dis_lr=1e-1,
adv_enc_lr=1e-1):
'''
model: a FewShotREModel instance
model_name: Name of the model
B: Batch size
N: Num of classes for each batch
K: Num of instances for each class in the support set
Q: Num of instances for each class in the query set
ckpt_dir: Directory of checkpoints
learning_rate: Initial learning rate
lr_step_size: Decay learning rate every lr_step_size steps
weight_decay: Rate of decaying weight
train_iter: Num of iterations of training
val_iter: Num of iterations of validating
val_step: Validate every val_step steps
test_iter: Num of iterations of testing
'''
print("Start training...")
# Init
if bert_optim:
print('Use bert optim!')
parameters_to_optimize = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
parameters_to_optimize = [
{'params': [p for n, p in parameters_to_optimize
if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in parameters_to_optimize
if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
optimizer = AdamW(parameters_to_optimize, lr=2e-5, correct_bias=False)
if self.adv:
optimizer_encoder = AdamW(parameters_to_optimize, lr=1e-5, correct_bias=False)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=warmup_step, t_total=train_iter)
else:
optimizer = pytorch_optim(model.parameters(),
learning_rate, weight_decay=weight_decay)
if self.adv:
optimizer_encoder = pytorch_optim(model.parameters(), lr=adv_enc_lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=lr_step_size)
if self.adv:
optimizer_dis = pytorch_optim(self.d.parameters(), lr=adv_dis_lr)
if load_ckpt:
state_dict = self.__load_model__(load_ckpt)['state_dict']
own_state = model.state_dict()
for name, param in state_dict.items():
if name not in own_state:
continue
own_state[name].copy_(param)
start_iter = 0
else:
start_iter = 0
if fp16:
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
model.train()
if self.adv:
self.d.train()
# Training
best_acc = 0
not_best_count = 0 # Stop training after several epochs without improvement.
iter_loss = 0.0
iter_loss_dis = 0.0
iter_right = 0.0
iter_right_dis = 0.0
iter_sample = 0.0
for it in range(start_iter, start_iter + train_iter):
if pair:
batch, label = next(self.train_data_loader)
if torch.cuda.is_available():
for k in batch:
batch[k] = batch[k].cuda()
label = label.cuda()
logits, pred = model(batch, N_for_train, K,
Q * N_for_train + na_rate * Q)
else:
support, query, label = next(self.train_data_loader)
if torch.cuda.is_available():
for k in support:
support[k] = support[k].cuda()
for k in query:
query[k] = query[k].cuda()
label = label.cuda()
logits, pred = model(support, query,
N_for_train, K, Q * N_for_train + na_rate * Q)
loss = model.loss(logits, label) / float(grad_iter)
right = model.accuracy(pred, label)
if fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
# torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 10)
else:
loss.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), 10)
if it % grad_iter == 0:
optimizer.step()
scheduler.step()
optimizer.zero_grad()
# Adv part
if self.adv:
support_adv = next(self.adv_data_loader)
if torch.cuda.is_available():
for k in support_adv:
support_adv[k] = support_adv[k].cuda()
features_ori = model.sentence_encoder(support)
features_adv = model.sentence_encoder(support_adv)
features = torch.cat([features_ori, features_adv], 0)
total = features.size(0)
dis_labels = torch.cat([torch.zeros((total//2)).long().cuda(),
torch.ones((total//2)).long().cuda()], 0)
dis_logits = self.d(features)
loss_dis = self.adv_cost(dis_logits, dis_labels)
_, pred = dis_logits.max(-1)
right_dis = float((pred == dis_labels).long().sum()) / float(total)
loss_dis.backward(retain_graph=True)
optimizer_dis.step()
optimizer_dis.zero_grad()
optimizer_encoder.zero_grad()
loss_encoder = self.adv_cost(dis_logits, 1 - dis_labels)
loss_encoder.backward(retain_graph=True)
optimizer_encoder.step()
optimizer_dis.zero_grad()
optimizer_encoder.zero_grad()
iter_loss_dis += self.item(loss_dis.data)
iter_right_dis += right_dis
iter_loss += self.item(loss.data)
iter_right += self.item(right.data)
iter_sample += 1
if self.adv:
sys.stdout.write('step: {0:4} | loss: {1:2.6f}, accuracy: {2:3.2f}%, dis_loss: {3:2.6f}, dis_acc: {4:2.6f}'
.format(it + 1, iter_loss / iter_sample,
100 * iter_right / iter_sample,
iter_loss_dis / iter_sample,
100 * iter_right_dis / iter_sample) +'\r')
else:
sys.stdout.write('step: {0:4} | loss: {1:2.6f}, accuracy: {2:3.2f}%'.format(it + 1, iter_loss / iter_sample, 100 * iter_right / iter_sample) +'\r')
sys.stdout.flush()
if (it + 1) % val_step == 0:
acc = self.eval(model, B, N_for_eval, K, Q, val_iter,
na_rate=na_rate, pair=pair)
model.train()
if acc > best_acc:
print('Best checkpoint')
torch.save({'state_dict': model.state_dict()}, save_ckpt)
best_acc = acc
iter_loss = 0.
iter_loss_dis = 0.
iter_right = 0.
iter_right_dis = 0.
iter_sample = 0.
print("\n####################\n")
print("Finish training " + model_name)
def main():
#os.environ["CUDA_VISIBLE_DEVICES"] = "0"
torch.set_num_threads(1)
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=False,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument("--task_name",
default=None,
type=str,
required=False,
help="The name of the task to train.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
## Other parameters
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 WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval or not.")
parser.add_argument("--eval_on",
default="dev",
help="Whether to run eval on the dev set or 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("--train_batch_size",
default=64,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=64,
type=int,
help="Total batch size for eval.")
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(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--weight_decay",
default=0.01,
type=float,
help="Weight deay 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("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
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 accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision 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(
'--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('--server_ip',
type=str,
default='',
help="Can be used for distant debugging.")
parser.add_argument('--server_port',
type=str,
default='',
help="Can be used for distant debugging.")
args = parser.parse_args()
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()
#processors = FormationProcessor
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:
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')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
processor = FormationProcessor()
tokenizer = BertTokenizer.from_pretrained(
'/home/ypd-19-2/SpERT/model/bertbase-20210122T060007Z-001/bertbase')
train_examples = None
num_train_optimization_steps = 0
if args.do_train:
train_examples = processor.get_train_examples()
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
if args.local_rank not in [-1, 0]:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
# Prepare model
config = BertConfig.from_pretrained(args.bert_model,
num_labels=1,
finetuning_task=args.task_name)
model = BertForSequenceClassification.from_pretrained(args.bert_model,
from_tf=False,
config=config)
if args.local_rank == 0:
torch.distributed.barrier(
) # Make sure only the first process in distributed training will download model & vocab
model.to(device)
param_optimizer = list(model.named_parameters())
no_decay = ['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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
warmup_steps = int(args.warmup_proportion * num_train_optimization_steps)
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
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)
if 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,
find_unused_parameters=True)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
#label_map = {i: label for i, label in enumerate(label_list, 1)}
if args.do_train:
train_features = convert_examples_to_features(train_examples,
args.max_seq_length,
tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.float)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)[0]
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:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)
#label_map = {i : label for i, label in enumerate(label_list,1)}
model_config = {
"bert_model": args.bert_model,
"do_lower": args.do_lower_case,
"max_seq_length": args.max_seq_length,
"num_labels": 1
}
json.dump(
model_config,
open(os.path.join(args.output_dir, "model_config.json"), "w"))
# Load a trained model and config that you have fine-tuned
else:
# Load a trained model and vocabulary that you have fine-tuned
model = BertForSequenceClassification.from_pretrained(args.output_dir)
tokenizer = BertTokenizer.from_pretrained(
args.output_dir, do_lower_case=args.do_lower_case)
model.to(device)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
loss_test = nn.L1Loss()
if args.eval_on == "dev":
eval_examples = processor.get_dev_examples()
elif args.eval_on == "test":
eval_examples = processor.get_test_examples()
else:
raise ValueError("eval on dev or test set only")
eval_features = convert_examples_to_features(eval_examples,
args.max_seq_length,
tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.float)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
y_true = []
y_pred = []
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)[0]
#logits = torch.argmax(F.log_softmax(logits,dim=2),dim=2)
input_mask = input_mask.to('cpu').numpy()
batch_loss = loss_test(logits, label_ids)
eval_loss += batch_loss
y_true.append(label_ids)
y_pred.append(logits)
print('eval_loss')
print(eval_loss / len(eval_dataloader))
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--data_dir",
default='/hdd/lujunyu/dataset/multi_turn_corpus/douban/',
type=str,
required=False,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument("--task_name",
default='ubuntu',
type=str,
required=False,
help="The name of the task to train.")
parser.add_argument(
"--output_dir",
default='/hdd/lujunyu/model/chatbert/douban_roberta_si_aug/',
type=str,
required=False,
help="The output directory where the model checkpoints will be written."
)
parser.add_argument(
"--dialog_augmentation_path",
default=
'/hdd/lujunyu/dataset/multi_turn_corpus/douban/train_augment_3.txt', ## train_augment_3.txt
type=str,
help="Whether to use augmentation")
## Other parameters
parser.add_argument(
"--init_model_path",
default='/hdd/lujunyu/dataset/bert/chinese_roberta_small_pytorch/',
type=str,
help="Initial checkpoint (usually from a pre-trained BERT model).")
parser.add_argument(
"--do_lower_case",
default=True,
action='store_true',
help=
"Whether to lower case the input text. True for uncased models, False for cased models."
)
parser.add_argument(
"--max_seq_length",
default=256,
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("--do_train",
default=True,
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_test",
default=True,
action='store_true',
help="Whether to run eval on the test set.")
parser.add_argument("--train_batch_size",
default=500,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=100,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-6,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=10.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_steps",
default=0.0,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--weight_decay",
default=1e-3,
type=float,
help="weight_decay")
parser.add_argument("--save_checkpoints_steps",
default=10000,
type=int,
help="How often to save the model checkpoint.")
parser.add_argument("--no_cuda",
default=False,
action='store_true',
help="Whether not to use CUDA when available")
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=10,
help=
"Number of updates steps to accumualte before performing a backward/update pass."
)
args = parser.parse_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.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = int(args.train_batch_size /
args.gradient_accumulation_steps)
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)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
bert_config = BertConfig.from_pretrained(args.init_model_path +
'config.json',
num_labels=2)
if args.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length {} because the BERT model was only trained up to sequence length {}"
.format(args.max_seq_length, bert_config.max_position_embeddings))
if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
if args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".
format(args.output_dir))
else:
os.makedirs(args.output_dir, exist_ok=True)
tokenizer = BertTokenizer.from_pretrained(args.init_model_path +
'vocab.txt',
do_lower_case=args.do_lower_case)
if args.dialog_augmentation_path:
train_dataset = DoubanDataset(file_path=args.dialog_augmentation_path,
max_seq_length=args.max_seq_length,
tokenizer=tokenizer)
else:
train_dataset = DoubanDataset(file_path=os.path.join(
args.data_dir, "train.txt"),
max_seq_length=args.max_seq_length,
tokenizer=tokenizer)
eval_dataset = DoubanDataset(file_path=os.path.join(
args.data_dir, "dev.txt"),
max_seq_length=args.max_seq_length,
tokenizer=tokenizer)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.train_batch_size,
sampler=RandomSampler(train_dataset),
num_workers=9)
eval_dataloader = torch.utils.data.DataLoader(
eval_dataset,
batch_size=args.eval_batch_size,
sampler=SequentialSampler(eval_dataset),
num_workers=9)
model = BertForSequenceClassification.from_pretrained(
'/hdd/lujunyu/model/chatbert/douban_roberta_si_aug_beifen/model.pt',
config=bert_config)
model.to(device)
num_train_steps = None
if args.do_train:
num_train_steps = int(
len(train_dataset) / args.train_batch_size /
args.gradient_accumulation_steps * args.num_train_epochs)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
# remove pooler, which is not used thus it produce None grad that break apex
param_optimizer = [n for n in param_optimizer]
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':
args.weight_decay
}, {
'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)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=num_train_steps)
else:
optimizer = None
scheduler = None
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)
global_step = 0
best_acc = 0.0
if args.do_train:
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_steps)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss, _ = model(input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
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
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step() # We have accumulated enought gradients
scheduler.step()
model.zero_grad()
global_step += 1
if (step + 1) % args.save_checkpoints_steps == 0:
### Evaluate at the end of epoches
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
logits_all = []
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss, logits = model(
input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
for i in range(len(logits)):
logits_all += [logits[i]]
tmp_eval_accuracy = accuracy(logits,
label_ids.reshape(-1))
eval_loss += tmp_eval_loss.mean().item()
eval_accuracy += tmp_eval_accuracy
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = eval_accuracy / nb_eval_examples
result = {
'eval_loss': eval_loss,
'eval_accuracy': eval_accuracy
}
output_eval_file = os.path.join(args.output_dir,
"eval_results_dev.txt")
with open(output_eval_file, "a") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
output_eval_file = os.path.join(args.output_dir,
"logits_dev.txt")
with open(output_eval_file, "w") as f:
for i in range(len(logits_all)):
for j in range(len(logits_all[i])):
f.write(str(logits_all[i][j]))
if j == len(logits_all[i]) - 1:
f.write("\n")
else:
f.write(" ")
### Save the best checkpoint
if best_acc < eval_accuracy:
try: ### Remove 'module' prefix when using DataParallel
state_dict = model.module.state_dict()
except AttributeError:
state_dict = model.state_dict()
torch.save(state_dict,
os.path.join(args.output_dir, "model.pt"))
best_acc = eval_accuracy
logger.info('Saving the best model in {}'.format(
os.path.join(args.output_dir, "model.pt")))
model.train()
def train(self):
if self.debug_mode: self.epochs = 2
print('加载dataloader')
# train_loader, valid_loader = self.create_dataloader()
train_dataloader, eval_dataloader, train_examples_length, valid_examples_length, eval_features = self.create_dataloader(
)
print('准备模型')
config = BertConfig.from_pretrained(self.model_name_or_path,
num_labels=2)
model = BertForSequenceClassification.from_pretrained(
self.model_name_or_path, self.args, config=config)
model.to(self.device)
print('准备优化器')
param_optimizer = list(model.named_parameters())
param_optimizer = [n for n in param_optimizer]
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':
self.weight_decay
}, {
'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=self.learning_rate,
eps=self.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=self.warmup_steps,
t_total=self.train_steps)
print('开始训练')
global_step = 0
best_acc = 0
train_loss = 0
num_train_optimization_steps = self.train_steps
for epoch in range(num_train_optimization_steps):
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(train_dataloader):
print('epoch:', epoch, 'batchIndex:', step)
batch = tuple(t.to(self.device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
loss.backward()
train_loss += loss.item()
scheduler.step()
optimizer.step()
optimizer.zero_grad()
global_step += 1
if (epoch + 1) % self.valid_step == 0:
logger.info("***** Report result *****")
logger.info(" %s = %s", 'global_step', str(global_step))
logger.info(" %s = %s", 'train loss', str(train_loss))
model.eval()
inference_labels = []
gold_labels = []
inference_logits = []
eval_loss, eval_accuracy = 0, 0
nb_eval_steps = 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(self.device)
input_mask = input_mask.to(self.device)
segment_ids = segment_ids.to(self.device)
label_ids = label_ids.to(self.device)
with torch.no_grad():
tmp_eval_loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
logits = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
inference_labels.append(np.argmax(logits, axis=1))
gold_labels.append(label_ids)
inference_logits.append(logits)
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
gold_labels = np.concatenate(gold_labels, 0)
inference_logits = np.concatenate(inference_logits, 0)
model.train()
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = self.accuracy(inference_logits, gold_labels)
result = {
'eval_loss': eval_loss,
'eval_F1': eval_accuracy,
'global_step': global_step,
'loss': train_loss
}
if eval_accuracy > best_acc:
logger.info("=" * 80)
logger.info("Best F1", eval_accuracy)
logger.info("Saving Model......")
best_acc = eval_accuracy
# Save a trained model
# model_to_save = model.module if hasattr(model,'module') else model # Only save the model it-self
# output_model_file = os.path.join(self.output_dir, "pytorch_model.bin")
# torch.save(model_to_save.state_dict(), self.output_model_file)
print("=" * 80)
else:
print("=" * 80)
model.train()
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 .tsv files (or other data files) for the task."
)
parser.add_argument("--ernie_model",
default="bert-base-cased",
type=str,
help="Ernie pre-trained model")
parser.add_argument("--embedding",
default="wikipedia2vec-base-cased",
type=str,
help="Embeddings")
parser.add_argument(
"--mapper",
default="wikipedia2vec-base-cased.bert-base-cased.linear",
type=str,
help="Embeddings")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
## Other parameters
parser.add_argument(
"--max_seq_length",
default=256,
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("--do_train",
default=False,
action='store_true',
help="Whether to run training.")
parser.add_argument("--ent",
required=True,
choices=("none", "concat", "replace"),
help="How to use entity embeddings.")
parser.add_argument("--do_eval",
default=False,
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--do_test",
default=False,
action='store_true',
help="Whether to run eval on the test set.")
parser.add_argument(
"--do_lower_case",
default=False,
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=4,
type=int,
help="Total batch size for evaluation.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=10.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
default=False,
action='store_true',
help="Whether not to use CUDA when available")
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 accumulate 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,
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('--threshold', type=float, default=.3)
args = parser.parse_args()
processors = FewrelProcessor
num_labels_task = 80
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:
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')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = int(args.train_batch_size /
args.gradient_accumulation_steps)
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)
if not args.do_train and not args.do_eval and not args.do_test:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
os.makedirs(args.output_dir, exist_ok=True)
processor = processors()
num_labels = num_labels_task
label_list = None
embedding = MappedEmbedding(load_embedding(args.embedding),
load_mapper(args.mapper))
model_embedding = load_embedding(args.ernie_model)
tokenizer = BertTokenizer.from_pretrained(args.ernie_model)
train_examples = None
num_train_steps = None
train_examples, label_list = processor.get_train_examples(args.data_dir)
label_list = sorted(label_list, key=lambda x: int(x[1:]))
num_train_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps * args.num_train_epochs)
model = EmbInputBertForSequenceClassification.from_pretrained(
args.ernie_model, num_labels=num_labels, output_attentions=True)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# 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
}]
t_total = num_train_steps
if args.local_rank != -1:
t_total = t_total // torch.distributed.get_world_size()
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_proportion *
t_total,
t_total=t_total)
global_step = 0
output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
patterns = ["# {name} #", "$ {name} $"]
if args.do_train:
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
tokenizer,
args.threshold,
patterns=patterns,
ent_type=args.ent,
embedding=embedding)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
output_loss_file = os.path.join(args.output_dir, "loss")
loss_fout = open(output_loss_file, 'w')
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
input_words = tokenizer.convert_ids_to_tokens(
batch[0].cpu().numpy().flatten())
input_vecs = [
embedding[w]
if w.startswith("ENTITY/") else model_embedding[w]
for w in input_words
]
input_vecs = np.array(input_vecs).reshape(batch[0].shape +
(-1, ))
batch[0] = torch.tensor(input_vecs)
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids,
attention_mask=input_mask,
token_type_ids=segment_ids,
labels=label_ids)[0]
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)
else:
loss.backward()
loss_fout.write("{}\n".format(loss.item()))
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
global_step += 1
model_to_save = model.module if hasattr(model, 'module') else model
output_model_file_step = os.path.join(
args.output_dir, "pytorch_model.bin_{}".format(global_step))
torch.save(model_to_save.state_dict(), output_model_file_step)
# Save a trained model
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
torch.save(model_to_save.state_dict(), output_model_file)
if args.do_eval or args.do_test:
del model
output_model_files = [
f for f in os.listdir(args.output_dir)
if f.startswith("pytorch_model.bin")
]
#output_model_files = ["pytorch_model.bin"] #TODO
for output_model_file in output_model_files:
model = EmbInputBertForSequenceClassification.from_pretrained(
args.ernie_model, num_labels=num_labels)
model.load_state_dict(
torch.load(os.path.join(args.output_dir, output_model_file)))
if args.fp16:
model.half()
model.to(device)
model.eval()
dsets = []
if args.do_eval:
dsets.append((processor.get_dev_examples, "eval"))
if args.do_test:
dsets.append((processor.get_test_examples, "test"))
for dset_func, dset_name in dsets:
features = convert_examples_to_features(dset_func(
args.data_dir),
label_list,
args.max_seq_length,
tokenizer,
args.threshold,
patterns=patterns,
ent_type=args.ent,
embedding=embedding)
step = output_model_file.replace("pytorch_model.bin", "")
fpred = open(
os.path.join(args.output_dir,
dset_name + f"_pred{step}.txt"), "w")
fgold = open(
os.path.join(args.output_dir,
dset_name + f"_gold{step}.txt"), "w")
fwords = open(
os.path.join(args.output_dir,
dset_name + f"_words{step}.txt"), "w")
all_input_ids = torch.tensor([f.input_ids for f in features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in features],
dtype=torch.long)
all_segment_ids = torch.tensor(
[f.segment_ids for f in features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in features],
dtype=torch.long)
data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
dataloader = DataLoader(data,
sampler=None,
batch_size=args.eval_batch_size)
acc = []
all_probs = []
for step, batch in enumerate(
tqdm(dataloader,
desc="Evaluation {} {}".format(
output_model_file, dset_name))):
input_words = tokenizer.convert_ids_to_tokens(
batch[0].cpu().numpy().flatten())
input_vecs = [
embedding[w]
if w.startswith("ENTITY/") else model_embedding[w]
for w in input_words
]
input_vecs = np.array(input_vecs).reshape(batch[0].shape +
(-1, ))
batch[0] = torch.tensor(input_vecs)
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
logits = model(input_ids,
attention_mask=input_mask,
token_type_ids=segment_ids)[0]
prob = torch.softmax(logits, -1)
all_probs.append(prob.detach().cpu().numpy())
predictions = prob.argmax(-1)
for a, b in zip(predictions, label_ids):
fgold.write("{}\n".format(label_list[b]))
fpred.write("{}\n".format(label_list[a]))
def train(train_dataset, model, tokenizer):
tb_writer = SummaryWriter()
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args['train_batch_size'])
t_total = len(train_dataloader) // args[
'gradient_accumulation_steps'] * args['num_train_epochs']
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_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
}]
warmup_steps = math.ceil(t_total * args['warmup_ratio'])
args['warmup_steps'] = warmup_steps if args['warmup_steps'] == 0 else args[
'warmup_steps']
optimizer = AdamW(optimizer_grouped_parameters,
lr=args['learning_rate'],
eps=args['adam_epsilon'],
betas=(0.9, 0.98))
#scheduler = WarmupConstantSchedule(optimizer, warmup_steps=args['warmup_steps'])#, t_total=t_total)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args['warmup_steps'],
t_total=t_total)
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'])
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args['num_train_epochs'])
logger.info(" Total train batch size = %d", args['train_batch_size'])
logger.info(" Gradient Accumulation steps = %d",
args['gradient_accumulation_steps'])
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_iterator = trange(int(args['num_train_epochs']), desc="Epoch")
for _ in train_iterator:
epoch_iterator = tqdm(train_dataloader, desc="Iteration")
for step, batch in enumerate(epoch_iterator):
model.train()
batch = tuple(t.to(device) for t in batch)
inputs = {
'input_ids':
batch[0],
'attention_mask':
batch[1],
'token_type_ids':
batch[2] if args['model_type'] in ['bert', 'xlnet'] else
None, # XLM don't use segment_ids
'labels':
batch[3]
}
outputs = model(**inputs)
loss = outputs[
0] # model outputs are always tuple in pytorch-transformers (see doc)
print("\r%f" % loss, end='')
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()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
args['max_grad_norm'])
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args['max_grad_norm'])
tr_loss += loss.item()
if (step + 1) % args['gradient_accumulation_steps'] == 0:
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:
# Log metrics
if args['evaluate_during_training']: # Only evaluate when single GPU otherwise metrics may not average well
results, _ = evaluate(model, tokenizer, evaluate=True)
for key, value in results.items():
tb_writer.add_scalar('eval_{}'.format(key), value,
global_step)
tb_writer.add_scalar('lr',
scheduler.get_lr()[0], global_step)
tb_writer.add_scalar('loss', (tr_loss - logging_loss) /
args['logging_steps'], global_step)
logging_loss = tr_loss
if 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)
logger.info("Saving model checkpoint to %s", output_dir)
return global_step, tr_loss / global_step
def main():
parser = argparse.ArgumentParser("")
parser.add_argument("--model", type=str, default='')
parser.add_argument("--resume", action='store_true')
parser.add_argument("--eval", action='store_true')
parser.add_argument("--batch_size", type=int, default=CFG.batch_size)
parser.add_argument("--nepochs", type=int, default=CFG.num_train_epochs)
parser.add_argument("--wsteps", type=int, default=CFG.warmup_steps)
parser.add_argument("--nlayers", type=int, default=CFG.num_hidden_layers)
parser.add_argument("--nahs", type=int, default=CFG.num_attention_heads)
parser.add_argument("--seed", type=int, default=7)
parser.add_argument("--lr", type=float, default=CFG.learning_rate)
parser.add_argument("--dropout", type=float, default=CFG.dropout)
parser.add_argument("--types", nargs='+', type=str,
default=['1JHC', '1JHN', '2JHC', '2JHH', '2JHN', '3JHC', '3JHH', '3JHN'],
help='3JHC,2JHC,1JHC,3JHH,2JHH,3JHN,2JHN,1JHN')
parser.add_argument("--train_file", default="train_mute_cp")
parser.add_argument("--test_file", default="test_mute_cp")
parser.add_argument("--pseudo_path", default="")
parser.add_argument("--pseudo", action='store_true')
parser.add_argument("--gen_pseudo", action='store_true')
parser.add_argument("--use_all", action='store_true')
parser.add_argument("--structure_file", default="structures_mu")
parser.add_argument("--contribution_file", default="scalar_coupling_contributions")
args = parser.parse_args()
print(args)
CFG.batch_size=args.batch_size
CFG.num_train_epochs=args.nepochs
CFG.warmup_steps=args.wsteps
CFG.num_hidden_layers=args.nlayers
CFG.num_attention_heads=args.nahs
CFG.learning_rate=args.lr
CFG.dropout=args.dropout
CFG.seed = args.seed
print(CFG.__dict__)
random.seed(CFG.seed)
np.random.seed(CFG.seed)
torch.manual_seed(CFG.seed)
#if not args.eval:
if True:
train_df = load_csv(args.train_file)
structures_df = load_csv(args.structure_file)
structures_df[['x', 'y', 'z']] -= structures_df.groupby('molecule_name')[['x', 'y', 'z']].transform('mean')
contributions_df = load_csv(args.contribution_file)
train_df = train_df.merge(contributions_df, how='left')
train_df = normalize_cols(train_df, ['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso'])
train_df = add_extra_features(train_df, structures_df)
train_df = train_df.fillna(1e08)
n_mols = train_df['molecule_name'].nunique()
train_df, valid_df = train_test_split(train_df, 5000 )
# only molecules with the args.types
print(train_df['molecule_name'].nunique())
mol_names_with_at = train_df[train_df['type'].isin(args.types)]['molecule_name'].unique()
train_df = train_df[train_df['molecule_name'].isin(mol_names_with_at)].reset_index(drop=True)
print(train_df['molecule_name'].nunique())
# Print the 5 rows of valid_df to verify whether the valid_df is the same as the previous experiment.
print(valid_df.head(5))
if args.pseudo:
test_df = load_csv(args.test_file)
logger.info(f'loading dataset - {args.pseudo_path} ...')
test_pseudo_df = pd.read_csv(args.pseudo_path)
#mol_names_jhn = train_df[test_df['type'].isin(['1JHN', '2JHN', '3JHN'])]['molecule_name'].unique()
#test_df = test_df[test_df['molecule_name'].isin(mol_names_jhn)].reset_index(drop=True)
test_df = add_extra_features(test_df, structures_df)
test_df = test_df.set_index('id')
test_pseudo_df = test_pseudo_df.set_index('id')
test_df[['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso']] = test_pseudo_df[['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso']]
test_df = test_df.reset_index()
#test_df = normalize_target(test_df)
test_df = normalize_cols(test_df, ['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso'])
#test_df = test_df.assign(fc=1e08, sd=1e08, pso=1e08, dso=1e08)
train_df['weight'] = 1.0
valid_df['weight'] = 1.0
test_df['weight'] = 1.0
n_mols = test_df['molecule_name'].nunique()
train_df = train_df.append(test_df).reset_index(drop=True)
else:
train_df['weight'] = 1.0
valid_df['weight'] = 1.0
if args.use_all:
train_df = train_df.append(valid_df)
print(f' n_train:{len(train_df)}, n_valid:{len(valid_df)}')
config = BertConfig(
3, # not used
hidden_size=CFG.hidden_size,
num_hidden_layers=CFG.num_hidden_layers,
num_attention_heads=CFG.num_attention_heads,
intermediate_size=CFG.intermediate_size,
hidden_dropout_prob=CFG.dropout,
attention_probs_dropout_prob=CFG.dropout,
)
model = cust_model.SelfAttn(config)
if args.model != "":
print("=> loading checkpoint '{}'".format(args.model))
checkpoint = torch.load(args.model)
CFG.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.model, checkpoint['epoch']))
model.cuda()
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
print('parameters: ', count_parameters(model))
n_gpu = torch.cuda.device_count()
if n_gpu > 1:
model = torch.nn.DataParallel(model)
# to produce the submission.csv
if args.eval:
test_df = load_csv(args.test_file)
structures_df = load_csv(args.structure_file)
structures_df[['x', 'y', 'z']] -= structures_df.groupby('molecule_name')[['x', 'y', 'z']].transform('mean')
test_df = add_extra_features(test_df, structures_df)
test_df = test_df.assign(fc=1e08, sd=1e08, pso=1e08, dso=1e08)
test_df['scalar_coupling_constant'] = 0
test_df['weight'] = 1.0
test_db = db.MolDB(test_df, CFG.max_seq_length)
test_loader = DataLoader(
test_db, batch_size=CFG.batch_size, shuffle=False,
num_workers=CFG.num_workers)
res_df = validate(test_loader, model, args.types)
res_df = unnormalize_cols(res_df, cols=['fc', 'sd', 'pso', 'dso'])
res_df = unnormalize_target(res_df, 'prediction1')
if args.gen_pseudo:
res_df['scalar_coupling_constant'] = res_df['prediction1']
res_df = res_df[res_df['id']>-1].sort_values('id')
res_df[['id', 'scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso']].to_csv(f'pseudo_{CFG.seed}.csv', index=False)
return
res_df['prediction4']= res_df[['fc', 'sd', 'pso', 'dso']].sum(1)
res_df['prediction']= res_df[['prediction1','prediction4']].mean(1)
res_df['scalar_coupling_constant'] = res_df['prediction']
res_df = res_df[res_df['id']>-1].sort_values('id')
os.makedirs('output', exist_ok=True)
res_df[['id', 'scalar_coupling_constant']].to_csv(f'output/submission_{CFG.seed}.csv', index=False)
return
train_db = db.MolDB(train_df, CFG.max_seq_length)
print('preloading dataset ...')
train_db = db.MolDB_FromDB(train_db, 10)
valid_db = db.MolDB(valid_df, CFG.max_seq_length)
num_train_optimization_steps = int(
len(train_db) / CFG.batch_size / CFG.gradient_accumulation_steps) * (CFG.num_train_epochs-CFG.start_epoch)
print('num_train_optimization_steps', num_train_optimization_steps)
train_loader = DataLoader(
train_db, batch_size=CFG.batch_size, shuffle=True,
num_workers=CFG.num_workers, pin_memory=True)
val_loader = DataLoader(
valid_db, batch_size=CFG.batch_size, shuffle=False,
num_workers=CFG.num_workers)
# 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=CFG.learning_rate,
weight_decay=CFG.weight_decay,
)
scheduler = WarmupLinearSchedule(optimizer, CFG.warmup_steps,
t_total=num_train_optimization_steps
)
def get_lr():
return scheduler.get_lr()[0]
if args.model != "":
if args.resume:
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
#for param_group in optimizer.param_groups:
# param_group['lr'] = CFG.learning_rate
mae_log_df = checkpoint['mae_log']
del checkpoint
else:
mae_log_df = pd.DataFrame(columns=(['EPOCH']+['LR']+args.types + ['OVERALL']) )
os.makedirs('log', exist_ok=True)
res_df = validate(val_loader, model, args.types)
res_df = unnormalize_cols(res_df, cols=['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso'])
res_df = unnormalize_target(res_df, 'prediction1')
res_df['prediction4']= res_df[['fc', 'sd', 'pso', 'dso']].sum(1)
res_df['prediction']= res_df[['prediction1','prediction4']].mean(1)
res_df.to_csv(f'log/valid_df_{"_".join(args.types)}.csv', index=False)
overall_mae, maes = metric(res_df, args.types)
print(overall_mae, maes)
curr_lr = get_lr()
print(f'initial learning rate:{curr_lr}')
for epoch in range(CFG.start_epoch, CFG.num_train_epochs):
# train for one epoch
#print(adjust_learning_rate(optimizer, epoch))
train(train_loader, model, optimizer, epoch, args.types, scheduler)
if epoch % CFG.test_freq == 0:
res_df = validate(val_loader, model, args.types)
res_df = unnormalize_cols(res_df, cols=['scalar_coupling_constant', 'fc', 'sd', 'pso', 'dso'])
res_df = unnormalize_target(res_df, 'prediction1')
res_df['prediction4']= res_df[['fc', 'sd', 'pso', 'dso']].sum(1)
res_df['prediction']= res_df[['prediction1','prediction4']].mean(1)
res_df.to_csv(f'log/valid_df_{"_".join(args.types)}.csv', index=False)
overall_mae, maes = metric(res_df, args.types)
# write log file
mae_row = dict([(typ, [mae]) for typ, mae in maes.items() if typ in args.types])
mae_row.update({'EPOCH':(epoch),'OVERALL':overall_mae, 'LR':curr_lr})
mae_log_df = mae_log_df.append(pd.DataFrame(mae_row), sort=False)
print(mae_log_df.tail(20))
mae_log_df.to_csv(f'log/{"_".join(args.types)}.csv', index=False)
#scheduler.step(overall_mae)
curr_lr = get_lr()
print(f'set the learning_rate: {curr_lr}')
# evaluate on validation set
batch_size = CFG.batch_size
pseudo_path = '' if not args.pseudo else '_' + args.pseudo_path
curr_model_name = (f'b{batch_size}_l{config.num_hidden_layers}_'
f'mh{config.num_attention_heads}_h{config.hidden_size}_'
f'd{CFG.dropout}_'
f'ep{epoch}_{"_".join(args.types)}_s{CFG.seed}{pseudo_path}.pt')
model_to_save = model.module if hasattr(model, 'module') else model # Only save the cust_model it-self
save_checkpoint({
'epoch': epoch + 1,
'arch': 'transformer',
'state_dict': model_to_save.state_dict(),
'mae_log': mae_log_df,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
},
FINETUNED_MODEL_PATH, curr_model_name
)
print('done')
'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=config["training"]["learning_rate"])
scheduler = WarmupLinearSchedule(
optimizer,
warmup_steps=config["training"]["warmup_proportion"],
t_total=config["training"]["total_training_steps"])
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer, FusedAdam
except:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=config["training"]["learning_rate"],
bias_correction=False)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
def train(train_dataloader, model, args, eval_dataloader=None):
"""
Trains the model. The params dict contains all parameters for training.
:param train_dataloader: torch.utils.data.dataloader.DataLoader
:param model: pytorch_transformers model
:param args: dict with the model parameters
Returns: model, train_loss_set, global_step, tr_loss / global_step
"""
tb_writer = SummaryWriter()
t_total = args["train_example_length"] // args["gradient_accumulation_steps"] * args["num_train_epochs"]
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_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}
]
warmup_steps = math.ceil(t_total * args["warmup_ratio"])
args["warmup_steps"] = warmup_steps if args["warmup_steps"] == 0 else args["warmup_steps"]
optimizer = AdamW(optimizer_grouped_parameters, lr=args["learning_rate"], eps=args["adam_epsilon"])
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args["warmup_steps"], t_total=t_total)
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'])
logger.info("***** Running training *****")
logger.info(" Num examples = %d", args['train_example_length'])
logger.info(" Num Epochs = %d", args['num_train_epochs'])
logger.info(" Total train batch size = %d", args['train_batch_size'])
logger.info(" Gradient Accumulation steps = %d", args['gradient_accumulation_steps'])
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_loss_set = []
for _ in trange(int(args["num_train_epochs"]), desc="Epoch"):
for step, batch in enumerate(tqdm_notebook(train_dataloader, desc="Iteration")):
model.train()
batch = tuple(t.to(args['device']) for t in batch)
inputs = {"input_ids": batch[0],
"attention_mask": batch[1],
"token_type_ids": batch[2] if args["model_type"] in ["bert", "xlnet"] else None,
# XLM don't use segment_ids
"labels": batch[3]}
outputs = model(**inputs)
loss = outputs[0] # model outputs are always tuple in pytorch-transformers (see doc)
train_loss_set.append(loss)
print("\r%f" % loss, end='')
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()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args["max_grad_norm"])
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args["max_grad_norm"])
tr_loss += loss.item()
if (step + 1) % args["gradient_accumulation_steps"] == 0:
scheduler.step() # Update learning rate schedule
optimizer.step()
model.zero_grad()
global_step += 1
if args['logging_steps'] > 0 and global_step % args["logging_steps"] == 0:
# Log metrics
if args["evaluate_during_training"]:
# Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(model, eval_dataloader, args)
# for key, value in results.items():
# tb_writer.add_scalar("eval_{}".format(key), value, global_step)
# tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
# tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args["logging_steps"], global_step)
# logging_loss = tr_loss
if 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)
logger.info("Saving model checkpoint to %s", output_dir)
return model, train_loss_set, global_step, tr_loss / global_step
def train(train_dataset,
model,
fp16,
n_gpus,
n_epochs,
learning_rate,
gradient_accumulation_steps,
per_gpu_train_batch_size,
device,
warmup_steps,
max_grad_norm,
callbacks,
nopbar=False,
layer_train_mask=None):
""" Train the model """
# tb_writer = SummaryWriter()
total_batchsize = per_gpu_train_batch_size * max(1, n_gpus)
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=total_batchsize)
total_steps = len(
train_dataloader) // gradient_accumulation_steps * n_epochs
def gen_trainable_params():
if layer_train_mask is None:
yield from model.named_parameters()
else:
for pname, pten in model.named_parameters():
match = re.findall('transformer\\.layer\\.([0-9]*)\\.', pname)
if match:
layer_num = int(match[0])
if layer_num in layer_train_mask:
yield pname, pten
else:
yield pname, pten
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in gen_trainable_params()
if not any(nd in n for nd in no_decay)
],
'weight_decay':
0.0
}, {
'params': [
p for n, p in gen_trainable_params()
if any(nd in n for nd in no_decay)
],
'weight_decay':
0.0
}]
optimizer = AdamW(optimizer_grouped_parameters,
lr=learning_rate,
eps=ADAM_EPSILON)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=total_steps)
if 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=FP16_OPT_LEVEL)
# multi-gpu training (should be after apex fp16 initialization)
if n_gpus > 1:
model = torch.nn.DataParallel(model)
# Train!
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", n_epochs)
logger.info(" Instantaneous batch size per GPU = %d",
per_gpu_train_batch_size)
logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d",
total_batchsize * gradient_accumulation_steps)
logger.info(" Gradient Accumulation steps = %d",
gradient_accumulation_steps)
logger.info(" Total optimization steps = %d", total_steps)
global_step = 0
accum_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_iterator = trange(int(n_epochs), desc="Epoch", disable=nopbar)
for _ in train_iterator:
epoch_iterator = tqdm(train_dataloader,
desc="Iteration",
disable=nopbar)
for step, batch in enumerate(epoch_iterator):
model.train()
batch = tuple(t.to(device) for t in batch)
# all_input_ids, all_input_mask, all_segment_ids, all_label_ids, spans
inputs = {
'input_ids': batch[0],
'attention_mask': batch[1],
'token_type_ids': batch[2],
'labels': batch[3],
'entstarts': batch[4].cpu(),
}
loss = model(**inputs)
if n_gpus > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel (not distributed) training
if gradient_accumulation_steps > 1:
loss = loss / gradient_accumulation_steps
if fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
max_grad_norm)
accum_loss += loss.item()
if (step + 1) % gradient_accumulation_steps == 0:
scheduler.step() # Update learning rate schedule
optimizer.step()
model.zero_grad()
global_step += 1
for callback_mod, cb in callbacks:
if global_step % callback_mod == 0:
cb(global_step, model, optimizer, scheduler,
accum_loss) # raises exceptions for break
return global_step, accum_loss / global_step
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model_type',
type=str,
default='openai-gpt',
help="model type: openai-gpt/gpt2/xlnet/...")
parser.add_argument('--model_name_or_path',
type=str,
default='openai-gpt',
help="pretrained model path")
parser.add_argument("--rel_lang",
action='store_true',
help="Use natural language to represent relations.")
parser.add_argument("--do_train", action='store_true', help="do training")
parser.add_argument("--do_eval", action='store_true', help="do evaluation")
parser.add_argument("--predict_rel",
action='store_true',
help="predict relation rather than objects")
parser.add_argument("--toy", action='store_true', help="test code")
parser.add_argument("--padding_text",
action='store_true',
help="xlnet needs a padding text")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument('--train_dataset',
type=str,
default='data/conceptnet/train100k.txt')
parser.add_argument('--eval_dataset1',
type=str,
default='data/conceptnet/dev1.txt')
parser.add_argument('--eval_dataset2',
type=str,
default='data/conceptnet/dev2.txt')
parser.add_argument('--eval_dataset', type=str, default='')
parser.add_argument('--test_dataset',
type=str,
default='data/conceptnet/test.txt')
parser.add_argument('--mask_parts',
type=str,
default='e1, r, or e2, which part to mask and predict')
parser.add_argument('--max_e1', type=int, default=10)
parser.add_argument('--max_r', type=int, default=5)
parser.add_argument('--max_e2', type=int, default=15)
parser.add_argument('--seed', type=int, default=123)
parser.add_argument("--no_pretrain",
action='store_true',
help="w/o pretrained parameters initialized")
parser.add_argument('--eval_per_steps', type=int, default=500)
parser.add_argument('--num_train_epochs', type=int, default=64)
parser.add_argument('--train_batch_size', type=int, default=64)
parser.add_argument('--eval_batch_size', type=int, default=16)
parser.add_argument('--max_grad_norm', type=int, default=1)
parser.add_argument('--learning_rate', type=float, default=1e-5)
parser.add_argument('--warmup_proportion', type=float, default=0.002)
parser.add_argument('--lr_schedule', type=str, default='warmup_linear')
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument('--adam_epsilon', type=float, default=1e-8)
args = parser.parse_args()
print(args)
set_seed(args.seed)
n_gpu = torch.cuda.device_count()
logger.info("device: {}, n_gpu {}".format(device, n_gpu))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# select model type
if args.model_type == "bert":
Tokenizer = BertTokenizer
Model = BertForMaskedLM
Config = BertConfig
else:
exit()
# load pretrained model
tokenizer = Tokenizer.from_pretrained(args.model_name_or_path)
if args.no_pretrain:
# from scratch
config = Config.from_pretrained(args.model_type)
model = Model(config)
else:
model = Model.from_pretrained(args.model_name_or_path)
print("vocab size:", len(tokenizer))
model.resize_token_embeddings(len(tokenizer))
model.to(device)
end_token = tokenizer.sep_token
mask_token_id = tokenizer.encode(tokenizer.mask_token)[0]
print("\nspecial tokens:", tokenizer.special_tokens_map)
# Load and encode the datasets
logger.info("Encoding dataset...")
train_dataset = load_comet_dataset(args.train_dataset,
end_token,
rel_lang=args.rel_lang,
toy=args.toy,
sep=True)
if args.eval_dataset:
eval_dataset = load_comet_dataset(args.eval_dataset,
end_token,
rel_lang=args.rel_lang,
toy=args.toy,
sep=True)
else:
eval_dataset1 = load_comet_dataset(args.eval_dataset1,
end_token,
rel_lang=args.rel_lang,
toy=args.toy,
sep=True)
eval_dataset2 = load_comet_dataset(args.eval_dataset2,
end_token,
rel_lang=args.rel_lang,
toy=args.toy,
sep=True)
eval_dataset = eval_dataset1 + eval_dataset2
test_dataset = load_comet_dataset(args.test_dataset,
end_token,
rel_lang=args.rel_lang,
toy=args.toy,
sep=True)
datasets = (train_dataset, eval_dataset, test_dataset)
encoded_datasets = tokenize_and_encode(datasets, tokenizer)
max_e1 = args.max_e1 + 1
max_r = args.max_r + 1
max_e2 = args.max_e2 + 1
input_length = max_e1 + max_r + max_e2
best_loss = 1e10
# Prepare inputs tensors and dataloaders
tensor_datasets = pre_process_datasets(encoded_datasets,
input_length,
max_e1,
max_r,
max_e2,
mask_parts=args.mask_parts,
mask_token=mask_token_id)
train_tensor_dataset, eval_tensor_dataset, test_tensor_dataset = tensor_datasets[
0], tensor_datasets[1], tensor_datasets[2]
train_data = TensorDataset(*train_tensor_dataset)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
eval_data = TensorDataset(*eval_tensor_dataset)
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
test_data = TensorDataset(*test_tensor_dataset)
test_sampler = SequentialSampler(test_data)
test_dataloader = DataLoader(test_data,
sampler=test_sampler,
batch_size=args.eval_batch_size)
# Prepare optimizer
if args.do_train:
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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
num_train_optimization_steps = len(
train_dataloader) * args.num_train_epochs
print("total steps:", num_train_optimization_steps)
num_warmup_steps = args.warmup_proportion * num_train_optimization_steps
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=num_warmup_steps,
t_total=num_train_optimization_steps)
if args.do_train:
nb_tr_steps = 0
model.train()
for cur_epoch_num in range(int(args.num_train_epochs)):
print("Epoch:", cur_epoch_num)
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)
input_ids, lm_labels, input_mask = batch
results = model(input_ids,
masked_lm_labels=lm_labels,
attention_mask=input_mask)
loss, logits = results
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
scheduler.step()
optimizer.step()
optimizer.zero_grad()
loss = loss.item()
nb_tr_steps += 1
if nb_tr_steps % (args.eval_per_steps / 10) == 0:
PPL = np.exp(loss) if loss < 300 else np.inf
print("Training loss:", loss, "ppl:", PPL)
if nb_tr_steps % args.eval_per_steps == 0:
model.eval()
# evaluate
eval_loss = evaluate(model, eval_dataloader, tokenizer,
max_e1, max_r, max_e2, args).item()
print("\n\nevaluating\neval loss:", eval_loss, "ppl",
np.exp(eval_loss) if eval_loss < 300 else np.inf)
# decide to save
if eval_loss < best_loss:
# save
save_model(model, tokenizer, args.output_dir)
print("model saved at step", nb_tr_steps)
print(str(datetime.datetime.now()))
print("prev loss:", best_loss, "cur loss:", eval_loss)
best_loss = eval_loss
# test
test_loss = evaluate(model, test_dataloader, tokenizer,
max_e1, max_r, max_e2, args).item()
print("\n\ntesting\ntest loss:", test_loss, "ppl:",
np.exp(test_loss) if test_loss < 300 else np.inf)
model.train()
if args.do_eval:
model.eval()
eval_loss = evaluate(model, eval_dataloader, tokenizer, max_e1, max_r,
max_e2, args).item()
print("\n\nevaluating\neval loss:", eval_loss, "ppl",
np.exp(eval_loss) if eval_loss < 300 else np.inf)
test_loss = evaluate(model, test_dataloader, tokenizer, max_e1, max_r,
max_e2, args).item()
print("\n\ntesting\ntest loss:", test_loss, "ppl:",
np.exp(test_loss) if test_loss < 300 else np.inf)
def main():
parser = argparse.ArgumentParser()
parser = add_xlmr_args(parser)
args = parser.parse_args()
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
processor = NerProcessor()
label_list = processor.get_labels()
num_labels = len(label_list) + 1 # add one for IGNORE label
train_examples = None
num_train_optimization_steps = 0
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
# preparing model configs
hidden_size = 768 if 'base' in args.pretrained_path else 1024 # TODO: move this inside model.__init__
device = 'cuda' if (torch.cuda.is_available()
and not args.no_cuda) else 'cpu'
print(device)
# creating model
model = XLMRForTokenClassification(pretrained_path=args.pretrained_path,
n_labels=num_labels,
hidden_size=hidden_size,
dropout_p=args.dropout,
device=device)
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
model.load_state_dict(torch.load('model_dir//model.pt'))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
model.to(device)
no_decay = ['bias', 'final_layer_norm.weight']
params = list(model.named_parameters())
optimizer_grouped_parameters = [{
'params':
[p for n, p in params if not any(nd in n for nd in no_decay)],
'weight_decay':
args.weight_decay
}, {
'params': [p for n, p in params if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
warmup_steps = int(args.warmup_proportion * num_train_optimization_steps)
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
# freeze model if necessary
if args.freeze_model:
logger.info("Freezing XLM-R model...")
for n, p in model.named_parameters():
if 'xlmr' in n and p.requires_grad:
p.requires_grad = False
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)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
label_map = {i: label for i, label in enumerate(label_list, 1)}
if args.do_train:
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
model.encode_word)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
train_data = create_dataset(train_features)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
# getting validation samples
val_examples = processor.get_dev_examples(args.data_dir)
val_features = convert_examples_to_features(val_examples, label_list,
args.max_seq_length,
model.encode_word)
val_data = create_dataset(val_features)
best_val_f1 = 0.0
for _ in tqdm(range(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
tbar = tqdm(train_dataloader, desc="Iteration")
model.train()
for step, batch in enumerate(tbar):
batch = tuple(t.to(device) for t in batch)
input_ids, label_ids, l_mask, valid_ids, = batch
loss = model(input_ids, label_ids, l_mask, valid_ids)
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()
torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
tbar.set_description('Loss = %.4f' % (tr_loss / (step + 1)))
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
logger.info("\nTesting on validation set...")
f1, report = evaluate_model(model, val_data, label_list,
args.eval_batch_size, device)
if f1 > best_val_f1:
best_val_f1 = f1
logger.info(
"\nFound better f1=%.4f on validation set. Saving model\n"
% (f1))
logger.info("%s\n" % (report))
torch.save(
model.state_dict(),
open(os.path.join(args.output_dir, 'model.pt'), 'wb'))
else:
logger.info("\nNo better F1 score: {}\n".format(f1))
else: # load a saved model
state_dict = torch.load(
open(os.path.join(args.output_dir, 'model.pt'), 'rb'))
model.load_state_dict(state_dict)
logger.info("Loaded saved model")
model.to(device)
if args.do_eval:
if args.eval_on == "dev":
eval_examples = processor.get_dev_examples(args.data_dir)
elif args.eval_on == "test":
eval_examples = processor.get_test_examples(args.data_dir)
else:
raise ValueError("eval on dev or test set only")
eval_features = convert_examples_to_features(eval_examples, label_list,
args.max_seq_length,
model.encode_word)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_data = create_dataset(eval_features)
f1_score, report = evaluate_model(model, eval_data, label_list,
args.eval_batch_size, device)
logger.info("\n%s", report)
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Writing results to file *****")
writer.write(report)
logger.info("Done.")
def train(self):
if not os.path.exists(self.output_dir):
os.makedirs(self.output_dir)
# logger.info(f'Fold {split_index + 1}')
train_dataloader, eval_dataloader, train_examples, eval_examples = self.create_dataloader(
)
num_train_optimization_steps = self.train_steps
# Prepare model
config = BertConfig.from_pretrained(self.model_name_or_path,
num_labels=self.num_labels)
model = BertForTokenClassification.from_pretrained(
self.model_name_or_path, self.args, config=config)
model.to(self.device)
model.train()
# Prepare optimizer
param_optimizer = list(model.named_parameters())
param_optimizer = [n for n in param_optimizer]
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':
self.weight_decay
}, {
'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=self.learning_rate,
eps=self.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=self.warmup_steps,
t_total=self.train_steps)
global_step = 0
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", self.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
best_acc = 0
best_MRR = 0
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
train_dataloader = cycle(train_dataloader)
for step in range(num_train_optimization_steps):
batch = next(train_dataloader)
batch = tuple(t.to(self.device) for t in batch)
input_ids, input_mask, segment_ids, label_domain, label_dependcy = batch
loss_domain, loss_dependcy = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
label_domain=label_domain,
label_dependcy=label_dependcy)
loss = loss_domain + loss_dependcy
tr_loss += loss.item()
train_loss = round(tr_loss / (nb_tr_steps + 1), 4)
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
loss.backward()
if (nb_tr_steps + 1) % self.gradient_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
scheduler.step()
global_step += 1
if (step + 1) % (self.eval_steps *
self.gradient_accumulation_steps) == 0:
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
logger.info("***** Report result *****")
logger.info(" %s = %s", 'global_step', str(global_step))
logger.info(" %s = %s", 'train loss', str(train_loss))
if self.do_eval and (step + 1) % (
self.eval_steps * self.gradient_accumulation_steps) == 0:
for file in ['dev.csv']:
inference_labels = []
gold_labels_domain = []
gold_labels_dependcy = []
inference_logits = []
scores_domain = []
scores_dependcy = []
ID = [x.guid for x in eval_examples]
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", self.eval_batch_size)
model.eval()
eval_loss_domain, eval_loss_dependcy, eval_accuracy_domain, eval_accuracy_dependcy = 0, 0, 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_domain, label_dependcy in eval_dataloader:
input_ids = input_ids.to(self.device)
input_mask = input_mask.to(self.device)
segment_ids = segment_ids.to(self.device)
label_domain = label_domain.to(self.device)
label_dependcy = label_dependcy.to(self.device)
with torch.no_grad():
batch_eval_loss_domain, batch_eval_loss_dependcy = model(
input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
label_domain=label_domain,
label_dependcy=label_dependcy)
logits_domain, logits_dependcy = model(
input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
logits_domain = logits_domain.view(
-1, self.num_labels).detach().cpu().numpy()
logits_dependcy = logits_dependcy.view(
-1, self.num_labels).detach().cpu().numpy()
label_domain = label_domain.view(-1).to('cpu').numpy()
label_dependcy = label_dependcy.view(-1).to(
'cpu').numpy()
scores_domain.append(logits_domain)
scores_dependcy.append(logits_dependcy)
gold_labels_domain.append(label_domain)
gold_labels_dependcy.append(label_dependcy)
eval_loss_domain += batch_eval_loss_domain.mean().item(
)
eval_loss_dependcy += batch_eval_loss_dependcy.mean(
).item()
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
gold_labels_domain = np.concatenate(gold_labels_domain, 0)
gold_labels_dependcy = np.concatenate(
gold_labels_dependcy, 0)
scores_domain = np.concatenate(scores_domain, 0)
scores_dependcy = np.concatenate(scores_dependcy, 0)
model.train()
eval_loss_domain = eval_loss_domain / nb_eval_steps
eval_loss_dependcy = eval_loss_dependcy / nb_eval_steps
eval_accuracy_domain = accuracyF1(scores_domain,
gold_labels_domain)
eval_accuracy_dependcy = accuracyF1(
scores_dependcy, gold_labels_dependcy)
print('eval_F1_domain', eval_accuracy_domain,
'eval_F1_dependcy', eval_accuracy_dependcy,
'global_step', global_step, 'loss', train_loss)
result = {
'eval_loss_domain': eval_loss_domain,
'eval_loss_dependcy': eval_loss_dependcy,
'eval_F1_domain': eval_accuracy_domain,
'eval_F1_dependcy': eval_accuracy_dependcy,
'global_step': global_step,
'loss': train_loss
}
output_eval_file = os.path.join(self.output_dir,
"eval_results.txt")
with open(output_eval_file, "a") as writer:
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
writer.write('*' * 80)
writer.write('\n')
if eval_accuracy_domain > best_acc:
print("=" * 80)
print("Best F1", eval_accuracy_domain)
print("Saving Model......")
# best_acc = eval_accuracy
best_acc = eval_accuracy_domain
# Save a trained model
model_to_save = model.module if hasattr(
model, 'module') else model
output_model_file = os.path.join(
self.output_dir, "pytorch_model.bin")
torch.save(model_to_save.state_dict(),
output_model_file)
print("=" * 80)
else:
print("=" * 80)
exit()
p for n, p in model.named_parameters()
if not any(nd in n for nd in no_decay)
],
'weight_decay':
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
}]
optimizer = AdamW(optimizer_grouped_parameters, lr=LR, eps=1e-8)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=WARMUP_STEPS,
t_total=len(train_dataloader) * EPOCHS)
# ## Learning
#
# Let's now define functions to train() and evaluate() the model:
def train(epoch, loss_vector=None, log_interval=200):
# Set model to training mode
model.train()
# Loop over each batch from the training set
for step, batch in enumerate(train_dataloader):
# Copy data to GPU if needed
def train(args, train_dataset, val_dataset, model, tokenizer):
""" Train the model """
pretrained_model = model[0]
adapter_model = model[1]
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size // args.gradient_accumulation_steps)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in adapter_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 adapter_model.named_parameters() 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)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
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.")
adapter_model, optimizer = amp.initialize(adapter_model, optimizer, opt_level=args.fp16_opt_level)
# multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1:
pretrained_model = torch.nn.DataParallel(pretrained_model)
adapter_model = torch.nn.DataParallel(adapter_model)
# Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1:
pretrained_model = torch.nn.parallel.DistributedDataParallel(pretrained_model, device_ids=[args.local_rank],
output_device=args.local_rank)
adapter_model = torch.nn.parallel.DistributedDataParallel(adapter_model, device_ids=[args.local_rank],
output_device=args.local_rank)
# Train!
logger.info("***** Running training *****")
logger.info(" Num train examples = %d", len(train_dataset)) #logging.info(f" Num train_examples = {len(train_examples)}")
logger.info(" Num Epochs = %d", args.num_train_epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
logger.info("Try resume from checkpoint")
if args.restore:
if os.path.exists(os.path.join(args.output_dir, 'global_step.bin')):
logger.info("Load last checkpoint data")
global_step = torch.load(os.path.join(args.output_dir, 'global_step.bin'))
output_dir = os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step))
logger.info("Load from output_dir {}".format(output_dir))
optimizer.load_state_dict(torch.load(os.path.join(output_dir, 'optimizer.bin')))
scheduler.load_state_dict(torch.load(os.path.join(output_dir, 'scheduler.bin')))
# args = torch.load(os.path.join(output_dir, 'training_args.bin'))
if hasattr(adapter_model, 'module'):
adapter_model.module.load_state_dict(torch.load(os.path.join(output_dir, 'pytorch_model.bin')))
else: # Take care of distributed/parallel training
adapter_model.load_state_dict(torch.load(os.path.join(output_dir, 'pytorch_model.bin')))
global_step += 1
start_epoch = int(global_step / len(train_dataloader))
start_step = global_step-start_epoch*len(train_dataloader)-1
logger.info("Start from global_step={} epoch={} step={}".format(global_step, start_epoch, start_step))
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter(log_dir="runs/" + args.my_model_name, purge_step=global_step)
else:
global_step = 0
start_epoch = 0
start_step = 0
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter(log_dir="runs/" + args.my_model_name, purge_step=global_step)
logger.info("Start from scratch")
else:
global_step = 0
start_epoch = 0
start_step = 0
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter(log_dir="runs/" + args.my_model_name, purge_step=global_step)
logger.info("Start from scratch")
tr_loss, logging_loss = 0.0, 0.0
pretrained_model.zero_grad()
adapter_model.zero_grad()
set_seed(args) # Added here for reproductibility (even between python 2 and 3)
for epoch in range(start_epoch, int(args.num_train_epochs)):
for step, batch in enumerate(train_dataloader):
start = time.time()
if args.restore and (step < start_step):
continue
# if args.restore and (flag_count < global_step):
# flag_count+=1
# continue
pretrained_model.eval()
adapter_model.train()
batch = tuple(t.to(args.device) for t in batch)
inputs = {'input_ids': batch[0],
'attention_mask': batch[1],
'token_type_ids': batch[2] if args.model_type in ['bert', 'xlnet'] else None, # XLM and RoBERTa don't use segment_ids
'labels': batch[3]}
pretrained_model_outputs = pretrained_model(**inputs)
outputs = adapter_model(pretrained_model_outputs,**inputs)
loss = outputs[0] # model outputs are always tuple in pytorch-transformers (see doc)
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
# epoch_iterator.set_description("loss {}".format(loss))
logger.info("Epoch {}/{} - Iter {} / {}, loss = {:.5f}, time used = {:.3f}s".format(epoch, int(args.num_train_epochs),step,
len(train_dataloader),
loss.item(),
time.time() - start))
if args.fp16:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(adapter_model.parameters(), args.max_grad_norm)
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
scheduler.step() # Update learning rate schedule
optimizer.step()
pretrained_model.zero_grad()
adapter_model.zero_grad()
global_step += 1
if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
# Log metrics
tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
tb_writer.add_scalar('loss', (tr_loss - logging_loss)/args.logging_steps, global_step)
logging_loss = tr_loss
if args.local_rank in [-1, 0] 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 = adapter_model.module if hasattr(adapter_model,
'module') else adapter_model # Take care of distributed/parallel training
model_to_save.save_pretrained(output_dir) # save to pytorch_model.bin model.state_dict()
torch.save(optimizer.state_dict(), os.path.join(output_dir, 'optimizer.bin'))
torch.save(scheduler.state_dict(), os.path.join(output_dir, 'scheduler.bin'))
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
torch.save(global_step, os.path.join(args.output_dir, 'global_step.bin'))
logger.info("Saving model checkpoint, optimizer, global_step to %s", output_dir)
if (global_step/args.save_steps) > args.max_save_checkpoints:
try:
shutil.rmtree(os.path.join(args.output_dir, 'checkpoint-{}'.format(global_step-args.max_save_checkpoints*args.save_steps)))
except OSError as e:
print(e)
if args.local_rank == -1 and args.evaluate_during_training and global_step %args.eval_steps== 0: # Only evaluate when single GPU otherwise metrics may not average well
model = (pretrained_model, adapter_model)
results = evaluate(args, val_dataset, model, tokenizer)
for key, value in results.items():
tb_writer.add_scalar('eval_{}'.format(key), value, global_step)
if args.max_steps > 0 and global_step > args.max_steps:
break
if args.max_steps > 0 and global_step > args.max_steps:
break
if args.local_rank in [-1, 0]:
tb_writer.close()
return global_step, tr_loss / global_step
def __init__(self,
params: dict,
dataloader: Dataset,
student: nn.Module,
teacher: nn.Module,
tokenizer: nn.Module):
logger.info('Initializing Distiller')
self.params = params
self.output_dir = params.output_dir
# self.multi_gpu = params.multi_gpu
self.student = student
self.teacher = teacher
self.tokenizer = tokenizer
self.dataloader = dataloader
# if self.params.n_gpu > 1:
# self.dataloader.split()
# self.num_steps_epoch = int(len(self.dataloader) / params.batch_size) + 1
self.num_steps_epoch = len(self.dataloader)
# print(len(self.dataloader), params.batch_size)
# print(self.num_steps_epoch, params.gradient_accumulation_steps, params.n_epoch)
# exit(0)
self.get_iterator()
self.temperature = params.temperature
assert self.temperature > 0.
self.alpha_ce = params.alpha_ce
self.alpha_mse = params.alpha_mse
self.alpha_cos = params.alpha_cos
assert self.alpha_ce >= 0.
assert self.alpha_mse >= 0.
assert self.alpha_cos >= 0.
assert self.alpha_ce + self.alpha_mse + self.alpha_cos > 0.
self.epoch = 0
self.n_iter = 0
self.n_total_iter = 0
self.n_sequences_epoch = 0
self.total_loss_epoch = 0
self.last_loss = 0
self.last_loss_ce = 0
if self.alpha_mse > 0.: self.last_loss_mse = 0
if self.alpha_cos > 0.: self.last_loss_cos = 0
self.last_log = 0
self.ce_loss_fct = nn.KLDivLoss(reduction='batchmean')
if self.alpha_mse > 0.:
self.mse_loss_fct = nn.MSELoss(reduction='sum')
if self.alpha_cos > 0.:
self.cosine_loss_fct = nn.CosineEmbeddingLoss(reduction='mean')
logger.info('--- Initializing model optimizer')
assert params.gradient_accumulation_steps >= 1
num_train_optimization_steps = int(self.num_steps_epoch / params.gradient_accumulation_steps * params.n_epoch) + 1
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in student.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad], 'weight_decay': params.weight_decay},
{'params': [p for n, p in student.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad], 'weight_decay': 0.0}
]
logger.info("------ Number of trainable parameters (student): %i" % sum([p.numel() for p in self.student.parameters() if p.requires_grad]))
logger.info("------ Number of parameters (student): %i" % sum([p.numel() for p in self.student.parameters()]))
self.optimizer = AdamW(optimizer_grouped_parameters,
lr=params.learning_rate,
eps=params.adam_epsilon,
betas=(0.9, 0.98))
warmup_steps = math.ceil(num_train_optimization_steps * params.warmup_prop)
self.scheduler = WarmupLinearSchedule(self.optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
# print("TOTAL", num_train_optimization_steps)
# print("WARM UP", warmup_steps)
# exit(0)
# if self.multi_gpu:
# if self.fp16:
# from apex.parallel import DistributedDataParallel
# logger.info("Using apex.parallel.DistributedDataParallel for distributed training.")
# self.student = DistributedDataParallel(self.student)
# else:
# from torch.nn.parallel import DistributedDataParallel
# logger.info("Using nn.parallel.DistributedDataParallel for distributed training.")
# self.student = DistributedDataParallel(self.student,
# device_ids=[params.local_rank],
# output_device=params.local_rank)
logger.info('--- Initializing Tensorboard')
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
logdir = os.path.join(self.params.output_dir, "tensorboard", current_time + '_' + socket.gethostname())
self.tensorboard = SummaryWriter(log_dir=logdir, flush_secs=60)
self.tensorboard.add_text(tag='config', text_string=str(self.params), global_step=0)
def train(self):
# Model
model = KobertCRF(config=self.model_config,
num_classes=len(self.tr_ds.ner_to_index))
model.to(self.device)
model.train()
# optim
train_examples_len = len(self.tr_ds)
val_examples_len = len(self.val_ds)
print("num of train: {}, num of val: {}".format(
train_examples_len, val_examples_len))
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
}]
# num_train_optimization_steps = int(train_examples_len / model_config.batch_size / model_config.gradient_accumulation_steps) * model_config.epochs
t_total = len(
self.tr_dl
) // self.model_config.gradient_accumulation_steps * self.model_config.epochs
optimizer = AdamW(optimizer_grouped_parameters,
lr=self.model_config.learning_rate,
eps=self.model_config.adam_epsilon)
scheduler = WarmupLinearSchedule(
optimizer,
warmup_steps=self.model_config.warmup_steps,
t_total=t_total)
n_gpu = torch.cuda.device_count()
# if n_gpu > 1:
# model = torch.nn.DataParallel(model)
# save
tb_writer = SummaryWriter('{}/runs'.format(self.model_dir))
checkpoint_manager = CheckpointManager(self.model_dir)
summary_manager = SummaryManager(self.model_dir)
# Train!
self.logger.info("***** Running training *****")
self.logger.info(" Num examples = %d", len(self.tr_ds))
self.logger.info(" Num Epochs = %d", self.model_config.epochs)
self.logger.info(" Instantaneous batch size per GPU = %d",
self.model_config.batch_size)
self.logger.info(" Gradient Accumulation steps = %d",
self.model_config.gradient_accumulation_steps)
self.logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
best_dev_acc, best_dev_loss = 0.0, 99999999999.0
best_steps = 0
model.zero_grad()
self.set_seed(
) # Added here for reproductibility (even between python 2 and 3)
# Train
train_iterator = trange(int(self.model_config.epochs), desc="Epoch")
for _epoch, _ in enumerate(train_iterator):
epoch_iterator = tqdm(self.tr_dl, desc="Iteration")
epoch = _epoch
for step, batch in enumerate(epoch_iterator):
model.train()
x_input, token_type_ids, y_real = map(
lambda elm: elm.to(self.device), batch)
log_likelihood, sequence_of_tags = model(
x_input, token_type_ids, y_real)
# loss: negative log-likelihood
loss = -1 * log_likelihood
if n_gpu > 1:
loss = loss.mean(
) # mean() to average on multi-gpu parallel training
if self.model_config.gradient_accumulation_steps > 1:
loss = loss / self.model_config.gradient_accumulation_steps
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
self.model_config.max_grad_norm)
tr_loss += loss.item()
if (step + 1
) % self.model_config.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
with torch.no_grad():
sequence_of_tags = torch.tensor(sequence_of_tags).to(
self.device)
print(sequence_of_tags.shape)
_tags = torch.squeeze(sequence_of_tags, dim=0)
mb_acc = (_tags == y_real).float()[
y_real != self.vocab.PAD_ID].mean()
tr_acc = mb_acc.item()
tr_loss_avg = tr_loss / global_step
tr_summary = {'loss': tr_loss_avg, 'acc': tr_acc}
# if step % 50 == 0:
print(
'epoch : {}, global_step : {}, tr_loss: {:.3f}, tr_acc: {:.2%}'
.format(epoch + 1, global_step, tr_summary['loss'],
tr_summary['acc']))
# training & evaluation log
if self.model_config.logging_steps > 0 and global_step % self.model_config.logging_steps == 0:
if self.model_config.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well
eval_summary, list_of_y_real, list_of_pred_tags = self.evaluate(
model, self.val_dl)
tb_writer.add_scalar('lr',
scheduler.get_lr()[0],
global_step)
tb_writer.add_scalars(
'loss', {
'train': (tr_loss - logging_loss) /
self.model_config.logging_steps,
'val':
eval_summary["eval_loss"]
}, global_step)
tb_writer.add_scalars(
'acc', {
'train': tr_acc,
'val': eval_summary["eval_acc"]
}, global_step)
print("eval acc: {}, loss: {}, global steps: {}".
format(eval_summary['eval_acc'],
eval_summary['eval_loss'],
global_step))
print("Average loss: {} at global step: {}".format(
(tr_loss - logging_loss) /
self.model_config.logging_steps, global_step))
logging_loss = tr_loss
# save model
if self.model_config.save_steps > 0 and global_step % self.model_config.save_steps == 0:
eval_summary, list_of_y_real, list_of_pred_tags = self.evaluate(
model, self.val_dl)
# Save model checkpoint
output_dir = os.path.join(self.model_config.output_dir,
'epoch-{}'.format(epoch + 1))
if not os.path.exists(output_dir):
os.makedirs(output_dir)
print("Saving model checkpoint to %s", output_dir)
state = {
'global_step': global_step + 1,
'model_state_dict': model.state_dict(),
'opt_state_dict': optimizer.state_dict()
}
summary = {'train': tr_summary, 'eval': eval_summary}
summary_manager.update(summary)
print("summary: ", summary)
summary_manager.save('summary.json')
# Save
is_best = eval_summary[
"eval_acc"] >= best_dev_acc # acc 기준 (원래는 train_acc가 아니라 val_acc로 해야)
if is_best:
best_dev_acc = eval_summary["eval_acc"]
best_dev_loss = eval_summary["eval_loss"]
best_steps = global_step
# if args.do_test:
# results_test = evaluate(model, test_dl, test=True)
# for key, value in results_test.items():
# tb_writer.add_scalar('test_{}'.format(key), value, global_step)
# logger.info("test acc: %s, loss: %s, global steps: %s", str(eval_summary['eval_acc']), str(eval_summary['eval_loss']), str(global_step))
checkpoint_manager.save_checkpoint(
state,
'best-epoch-{}-step-{}-acc-{:.3f}.bin'.format(
epoch + 1, global_step, best_dev_acc))
print(
"Saving model checkpoint as best-epoch-{}-step-{}-acc-{:.3f}.bin"
.format(epoch + 1, global_step, best_dev_acc))
# print classification report and save confusion matrix
cr_save_path = self.model_dir + '/best-epoch-{}-step-{}-acc-{:.3f}-cr.csv'.format(
epoch + 1, global_step, best_dev_acc)
cm_save_path = self.model_dir + '/best-epoch-{}-step-{}-acc-{:.3f}-cm.png'.format(
epoch + 1, global_step, best_dev_acc)
self.save_cr_and_cm(list_of_y_real,
list_of_pred_tags,
cr_save_path=cr_save_path,
cm_save_path=cm_save_path)
else:
torch.save(
state,
os.path.join(
output_dir,
'model-epoch-{}-step-{}-acc-{:.3f}.bin'.
format(epoch + 1, global_step,
eval_summary["eval_acc"])))
print(
"Saving model checkpoint as model-epoch-{}-step-{}-acc-{:.3f}.bin"
.format(epoch + 1, global_step,
eval_summary["eval_acc"]))
tb_writer.close()
print("global_step = {}, average loss = {}".format(
global_step, tr_loss / global_step))
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
warmup_steps = int(args.warmup_proportion * num_train_optimization_steps)
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
# For our experiment, the following can be ignored
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)
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 .tsv files (or other data files) for the task."
)
parser.add_argument("--model_type",
default=None,
type=str,
required=True,
help="Model type selected in the list: " +
", ".join(MODEL_CLASSES.keys()))
parser.add_argument(
"--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model or shortcut name selected in the list: "
+ ", ".join(ALL_MODELS))
parser.add_argument(
"--meta_path",
default=None,
type=str,
required=False,
help="Path to pre-trained model or shortcut name selected in the list: "
+ ", ".join(ALL_MODELS))
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and 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_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_test",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--evaluate_during_training",
action='store_true',
help="Rul evaluation during training at each logging step.")
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 deay 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("--eval_steps", default=-1, type=int, help="")
parser.add_argument("--lstm_hidden_size", default=300, type=int, help="")
parser.add_argument("--lstm_layers", default=2, type=int, help="")
parser.add_argument("--lstm_dropout", default=0.5, type=float, help="")
parser.add_argument("--train_steps", default=-1, type=int, help="")
parser.add_argument("--report_steps", default=-1, type=int, help="")
parser.add_argument("--warmup_steps",
default=0,
type=int,
help="Linear warmup over warmup_steps.")
parser.add_argument("--split_num", default=3, type=int, help="text split")
parser.add_argument('--logging_steps',
type=int,
default=50,
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('--device_id',
type=str,
default='0',
help="The CUDA device is for training.")
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.device_id
# 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)
try:
os.makedirs(args.output_dir)
except:
pass
tokenizer = BertTokenizer.from_pretrained(args.model_name_or_path,
do_lower_case=args.do_lower_case)
config = BertConfig.from_pretrained(args.model_name_or_path, num_labels=2)
# Prepare model
model = BertForSequenceClassification.from_pretrained(
args.model_name_or_path, args, config=config)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif args.n_gpu > 1:
model = torch.nn.DataParallel(model)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
if args.do_train:
# Prepare data loader
train_examples = read_examples(os.path.join(args.data_dir,
'train.csv'),
is_training=True)
train_features = convert_examples_to_features(train_examples,
tokenizer,
args.max_seq_length,
args.split_num, True)
all_input_ids = torch.tensor(select_field(train_features, 'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(train_features,
'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(train_features,
'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size //
args.gradient_accumulation_steps)
num_train_optimization_steps = args.train_steps
# Prepare optimizer
param_optimizer = list(model.named_parameters())
# hack to remove pooler, which is not used
# thus it produce None grad that break apex
param_optimizer = [n for n in param_optimizer]
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':
args.weight_decay
}, {
'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)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=args.train_steps)
global_step = 0
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
best_acc = 0
model.train()
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
bar = tqdm(range(num_train_optimization_steps),
total=num_train_optimization_steps)
train_dataloader = cycle(train_dataloader)
for step in bar:
batch = next(train_dataloader)
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
if args.n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.fp16 and args.loss_scale != 1.0:
loss = loss * args.loss_scale
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
tr_loss += loss.item()
train_loss = round(
tr_loss * args.gradient_accumulation_steps / (nb_tr_steps + 1),
4)
bar.set_description("loss {}".format(train_loss))
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
if (nb_tr_steps + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear.get_lr(
global_step, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
scheduler.step()
optimizer.step()
optimizer.zero_grad()
global_step += 1
if (step + 1) % (args.eval_steps *
args.gradient_accumulation_steps) == 0:
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
logger.info("***** Report result *****")
logger.info(" %s = %s", 'global_step', str(global_step))
logger.info(" %s = %s", 'train loss', str(train_loss))
if args.do_eval and (step + 1) % (
args.eval_steps * args.gradient_accumulation_steps) == 0:
for file in ['dev.csv']:
inference_labels = []
gold_labels = []
inference_logits = []
eval_examples = read_examples(os.path.join(
args.data_dir, file),
is_training=True)
eval_features = convert_examples_to_features(
eval_examples, tokenizer, args.max_seq_length,
args.split_num, False)
all_input_ids = torch.tensor(select_field(
eval_features, 'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(
eval_features, 'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(
eval_features, 'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(
eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
logits = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
inference_labels.append(np.argmax(logits, axis=1))
gold_labels.append(label_ids)
inference_logits.append(logits)
eval_loss += tmp_eval_loss.mean().item()
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
gold_labels = np.concatenate(gold_labels, 0)
inference_logits = np.concatenate(inference_logits, 0)
model.train()
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = accuracy(inference_logits, gold_labels)
result = {
'eval_loss': eval_loss,
'eval_F1': eval_accuracy,
'global_step': global_step,
'loss': train_loss
}
output_eval_file = os.path.join(args.output_dir,
"eval_results.txt")
with open(output_eval_file, "a") as writer:
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
writer.write('*' * 80)
writer.write('\n')
if eval_accuracy > best_acc and 'dev' in file:
print("=" * 80)
print("Best F1", eval_accuracy)
print("Saving Model......")
best_acc = eval_accuracy
# Save a trained model
model_to_save = model.module if hasattr(
model,
'module') else model # Only save the model it-self
output_model_file = os.path.join(
args.output_dir, "pytorch_model.bin")
torch.save(model_to_save.state_dict(),
output_model_file)
print("=" * 80)
else:
print("=" * 80)
if args.do_test:
del model
gc.collect()
args.do_train = False
model = BertForSequenceClassification.from_pretrained(os.path.join(
args.output_dir, "pytorch_model.bin"),
args,
config=config)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif args.n_gpu > 1:
model = torch.nn.DataParallel(model)
for file, flag in [('dev.csv', 'dev'), ('test.csv', 'test')]:
inference_labels = []
gold_labels = []
eval_examples = read_examples(os.path.join(args.data_dir, file),
is_training=False)
eval_features = convert_examples_to_features(
eval_examples, tokenizer, args.max_seq_length, args.split_num,
False)
all_input_ids = torch.tensor(select_field(eval_features,
'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(eval_features,
'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(
eval_features, 'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(
input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask).detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
inference_labels.append(logits)
gold_labels.append(label_ids)
gold_labels = np.concatenate(gold_labels, 0)
logits = np.concatenate(inference_labels, 0)
print(flag, accuracy(logits, gold_labels))
if flag == 'test':
df = pd.read_csv(os.path.join(args.data_dir, file))
df['label_0'] = logits[:, 0]
df['label_1'] = logits[:, 1]
#df['label_2']=logits[:,2]
df[['id', 'entity', 'label_0',
'label_1']].to_csv(os.path.join(args.output_dir,
"test_pb.csv"),
index=False)
def active_train(data_loader=None,
model=None,
Epochs=5,
soft_loader=None,
args=None):
config = BertConfig.from_pretrained(args.bert_model,
num_labels=num_labels,
finetuning_task=args.task_name)
if model == None:
model = Ner.from_pretrained(args.bert_model,
from_tf=False,
config=config)
return_model = Ner.from_pretrained(args.bert_model,
from_tf=False,
config=config)
model.to(device)
return_model.to(device)
param_optimizer = list(model.named_parameters())
no_decay = ['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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
num_train_optimization_steps = int(
len(data_loader.dataset) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs #2190
warmup_steps = int(args.warmup_proportion * num_train_optimization_steps)
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
current_train_size = 0
if soft_loader:
current_train_size = len(data_loader.dataset) + len(
soft_loader.dataset)
else:
current_train_size = len(data_loader.dataset)
print('Training on {} data'.format(current_train_size))
model.train()
tr_loss = 2020
for epoch_idx in trange(int(Epochs), desc="Epoch"):
current_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(tqdm(data_loader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids, valid_ids, l_mask = batch
loss = model(input_ids, segment_ids, input_mask, label_ids,
valid_ids, l_mask)
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:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer),
args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
current_loss += loss
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
if soft_loader:
for input_ids, input_mask, segment_ids, soft_labels, valid_ids, l_mask in tqdm(
soft_loader, desc="Soft Training"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
soft_labels = soft_labels.to(device)
l_mask = l_mask.to(device)
#with torch.no_grad():
logits = model(input_ids,
segment_ids,
input_mask,
valid_ids=valid_ids,
attention_mask_label=l_mask)
#logits = F.softmax(logits, dim=2)
logits = logits.detach().cpu().float()
soft_labels = soft_labels.detach().cpu().float()
pos_weight = torch.ones([num_labels
]) # All weights are equal to 1
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
loss = 0
for i in range(len(logits)):
turncate_len = np.count_nonzero(
l_mask[i].detach().cpu().numpy())
logit = logits[i][:turncate_len]
soft_label = soft_labels[i][:turncate_len]
loss += criterion(logit, soft_label)
loss = Variable(loss, requires_grad=True)
current_loss += loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(),
args.max_grad_norm)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
if current_loss <= tr_loss:
return_model.load_state_dict(model.state_dict())
tr_loss = current_loss
return return_model
def __init__(self, params: dict, dataloader: Dataset,
token_probs: torch.tensor, student: nn.Module,
teacher: nn.Module):
logger.info('Initializing Distiller')
self.params = params
self.dump_path = params.dump_path
self.multi_gpu = params.multi_gpu
self.fp16 = params.fp16
self.student = student
self.teacher = teacher
self.dataloader = dataloader
if self.params.n_gpu > 1:
self.dataloader.split()
self.get_iterator(seed=params.seed)
self.temperature = params.temperature
assert self.temperature > 0.
self.alpha_ce = params.alpha_ce
self.alpha_mlm = params.alpha_mlm
self.alpha_mse = params.alpha_mse
assert self.alpha_ce >= 0.
assert self.alpha_mlm >= 0.
assert self.alpha_mse >= 0.
assert self.alpha_ce + self.alpha_mlm + self.alpha_mse > 0.
self.mlm_mask_prop = params.mlm_mask_prop
assert 0.0 <= self.mlm_mask_prop <= 1.0
assert params.word_mask + params.word_keep + params.word_rand == 1.0
self.pred_probs = torch.FloatTensor(
[params.word_mask, params.word_keep, params.word_rand])
self.pred_probs = self.pred_probs.to(
f'cuda:{params.local_rank}'
) if params.n_gpu > 0 else self.pred_probs
self.token_probs = token_probs.to(
f'cuda:{params.local_rank}') if params.n_gpu > 0 else token_probs
if self.fp16:
self.pred_probs = self.pred_probs.half()
self.token_probs = self.token_probs.half()
self.epoch = 0
self.n_iter = 0
self.n_total_iter = 0
self.n_sequences_epoch = 0
self.total_loss_epoch = 0
self.last_loss = 0
self.last_loss_ce = 0
self.last_loss_mlm = 0
self.last_loss_mse = 0
self.ce_loss_fct = nn.KLDivLoss(reduction='batchmean')
self.mlm_loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
self.mse_loss_fct = nn.MSELoss(reduction='sum')
logger.info('--- Initializing model optimizer')
assert params.gradient_accumulation_steps >= 1
self.num_steps_epoch = int(
len(self.dataloader) / params.batch_size) + 1
num_train_optimization_steps = int(
self.num_steps_epoch / params.gradient_accumulation_steps *
params.n_epoch) + 1
warmup_steps = math.ceil(num_train_optimization_steps *
params.warmup_prop)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in student.named_parameters()
if not any(nd in n for nd in no_decay) and p.requires_grad
],
'weight_decay':
params.weight_decay
}, {
'params': [
p for n, p in student.named_parameters()
if any(nd in n for nd in no_decay) and p.requires_grad
],
'weight_decay':
0.0
}]
logger.info(
"------ Number of trainable parameters (student): %i" % sum([
p.numel() for p in self.student.parameters() if p.requires_grad
]))
logger.info("------ Number of parameters (student): %i" %
sum([p.numel() for p in self.student.parameters()]))
self.optimizer = AdamW(optimizer_grouped_parameters,
lr=params.learning_rate,
eps=params.adam_epsilon,
betas=(0.9, 0.98))
self.scheduler = WarmupLinearSchedule(
self.optimizer,
warmup_steps=warmup_steps,
t_total=num_train_optimization_steps)
if self.fp16:
try:
from apex import amp
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
)
logger.info(
f"Using fp16 training: {self.params.fp16_opt_level} level")
self.student, self.optimizer = amp.initialize(
self.student,
self.optimizer,
opt_level=self.params.fp16_opt_level)
self.teacher = self.teacher.half()
if self.multi_gpu:
if self.fp16:
from apex.parallel import DistributedDataParallel
logger.info(
"Using apex.parallel.DistributedDataParallel for distributed training."
)
self.student = DistributedDataParallel(self.student)
else:
from torch.nn.parallel import DistributedDataParallel
logger.info(
"Using nn.parallel.DistributedDataParallel for distributed training."
)
self.student = DistributedDataParallel(
self.student,
device_ids=[params.local_rank],
output_device=params.local_rank)
self.is_master = params.is_master
if self.is_master:
logger.info('--- Initializing Tensorboard')
self.tensorboard = SummaryWriter(
log_dir=os.path.join(self.dump_path, 'log', 'train'))
self.tensorboard.add_text(tag='config',
text_string=str(self.params),
global_step=0)
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 .tsv files (or other data files) for the task."
)
parser.add_argument("--model_type",
default=None,
type=str,
required=True,
help="Model type selected in the list: " +
", ".join(MODEL_CLASSES.keys()))
parser.add_argument(
"--model_name_or_path",
default=None,
type=str,
required=True,
help="Path to pre-trained model or shortcut name selected in the list: "
+ ", ".join(ALL_MODELS))
parser.add_argument(
"--meta_path",
default=None,
type=str,
required=False,
help="Path to pre-trained model or shortcut name selected in the list: "
+ ", ".join(ALL_MODELS))
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
parser.add_argument(
'--classifier',
default='guoday',
type=str,
required=True,
help='classifier type, guoday or MLP or GRU_MLP or ...')
parser.add_argument('--optimizer',
default='RAdam',
type=str,
required=True,
help='optimizer we use, RAdam or ...')
parser.add_argument("--do_label_smoothing",
default='yes',
type=str,
required=True,
help="Whether to do label smoothing. yes or no.")
parser.add_argument('--draw_loss_steps',
default=1,
type=int,
required=True,
help='training steps to draw loss')
parser.add_argument('--label_name',
default='label',
type=str,
required=True,
help='label name in original train set index')
## 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_train",
default='yes',
type=str,
required=True,
help="Whether to run training. yes or no.")
parser.add_argument("--do_test",
default='yes',
type=str,
required=True,
help="Whether to run training. yes or no.")
parser.add_argument("--do_eval",
default='yes',
type=str,
required=True,
help="Whether to run eval on the dev set. yes or no.")
parser.add_argument(
"--evaluate_during_training",
action='store_true',
help="Rul evaluation during training at each logging step.")
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 deay 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("--eval_steps", default=200, type=int, help="")
parser.add_argument("--lstm_hidden_size", default=300, type=int, help="")
parser.add_argument("--lstm_layers", default=2, type=int, help="")
parser.add_argument("--dropout", default=0.5, type=float, help="")
parser.add_argument("--train_steps", default=-1, type=int, help="")
parser.add_argument("--report_steps", default=-1, type=int, help="")
parser.add_argument("--warmup_steps",
default=0,
type=int,
help="Linear warmup over warmup_steps.")
parser.add_argument("--split_num", default=3, type=int, help="text split")
parser.add_argument('--logging_steps',
type=int,
default=50,
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.")
args = parser.parse_args()
# 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)
try:
os.makedirs(args.output_dir)
except:
pass
tokenizer = BertTokenizer.from_pretrained(args.model_name_or_path,
do_lower_case=args.do_lower_case)
# tensorboard_log_dir = args.output_dir
# loss_now = tf.placeholder(dtype=tf.float32, name='loss_now')
# loss_mean = tf.placeholder(dtype=tf.float32, name='loss_mean')
# loss_now_variable = loss_now
# loss_mean_variable = loss_mean
# train_loss = tf.summary.scalar('train_loss', loss_now_variable)
# dev_loss_mean = tf.summary.scalar('dev_loss_mean', loss_mean_variable)
# merged = tf.summary.merge([train_loss, dev_loss_mean])
config = BertConfig.from_pretrained(args.model_name_or_path, num_labels=3)
config.hidden_dropout_prob = args.dropout
# Prepare model
if args.do_train == 'yes':
model = BertForSequenceClassification.from_pretrained(
args.model_name_or_path, args, config=config)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif args.n_gpu > 1:
model = torch.nn.DataParallel(model)
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
if args.do_train == 'yes':
print(
'________________________now training______________________________'
)
# Prepare data loader
train_examples = read_examples(os.path.join(args.data_dir,
'train.csv'),
is_training=True,
label_name=args.label_name)
train_features = convert_examples_to_features(train_examples,
tokenizer,
args.max_seq_length,
args.split_num, True)
# print('train_feature_size=', train_features.__sizeof__())
all_input_ids = torch.tensor(select_field(train_features, 'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(train_features,
'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(train_features,
'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
# print('train_data=',train_data[0])
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size //
args.gradient_accumulation_steps)
num_train_optimization_steps = args.train_steps
# Prepare optimizer
param_optimizer = list(model.named_parameters())
# hack to remove pooler, which is not used
# thus it produce None grad that break apex
param_optimizer = [n for n in param_optimizer]
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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.optimizer == 'RAdam':
optimizer = RAdam(optimizer_grouped_parameters,
lr=args.learning_rate)
else:
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=args.train_steps)
global_step = 0
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
best_acc = 0
model.train()
tr_loss = 0
loss_batch = 0
nb_tr_examples, nb_tr_steps = 0, 0
bar = tqdm(range(num_train_optimization_steps),
total=num_train_optimization_steps)
train_dataloader = cycle(train_dataloader)
# with tf.Session() as sess:
# summary_writer = tf.summary.FileWriter(tensorboard_log_dir, sess.graph)
# sess.run(tf.global_variables_initializer())
list_loss_mean = []
bx = []
eval_F1 = []
ax = []
for step in bar:
batch = next(train_dataloader)
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
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
tr_loss += loss.item()
loss_batch += loss.item()
train_loss = round(
tr_loss * args.gradient_accumulation_steps / (nb_tr_steps + 1),
4)
bar.set_description("loss {}".format(train_loss))
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if args.fp16:
# optimizer.backward(loss)
loss.backward()
else:
loss.backward()
# draw loss every n docs
if (step + 1) % int(args.draw_loss_steps /
(args.train_batch_size /
args.gradient_accumulation_steps)) == 0:
list_loss_mean.append(round(loss_batch, 4))
bx.append(step + 1)
plt.plot(bx,
list_loss_mean,
label='loss_mean',
linewidth=1,
color='b',
marker='o',
markerfacecolor='green',
markersize=2)
plt.savefig(args.output_dir + '/labeled.jpg')
loss_batch = 0
# paras update every batch data.
if (nb_tr_steps + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear.get_lr(
global_step, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
scheduler.step()
optimizer.step()
optimizer.zero_grad()
global_step += 1
# report results every 200 real batch.
if step % (args.eval_steps *
args.gradient_accumulation_steps) == 0 and step > 0:
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
logger.info("***** Report result *****")
logger.info(" %s = %s", 'global_step', str(global_step))
logger.info(" %s = %s", 'train loss', str(train_loss))
# do evaluation totally 10 times during training stage.
if args.do_eval == 'yes' and (step + 1) % int(
num_train_optimization_steps / 10) == 0 and step > 450:
for file in ['dev.csv']:
inference_labels = []
gold_labels = []
inference_logits = []
eval_examples = read_examples(os.path.join(
args.data_dir, file),
is_training=True,
label_name=args.label_name)
eval_features = convert_examples_to_features(
eval_examples, tokenizer, args.max_seq_length,
args.split_num, False)
all_input_ids = torch.tensor(select_field(
eval_features, 'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(
eval_features, 'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(
eval_features, 'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(
eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask,
labels=label_ids)
logits = model(input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
inference_labels.append(np.argmax(logits, axis=1))
gold_labels.append(label_ids)
inference_logits.append(logits)
eval_loss += tmp_eval_loss.mean().item()
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
gold_labels = np.concatenate(gold_labels, 0)
inference_labels = np.concatenate(inference_labels, 0)
inference_logits = np.concatenate(inference_logits, 0)
model.train()
###############################################
num_gold_0 = np.sum(gold_labels == 0)
num_gold_1 = np.sum(gold_labels == 1)
num_gold_2 = np.sum(gold_labels == 2)
right_0 = 0
right_1 = 0
right_2 = 0
error_0 = 0
error_1 = 0
error_2 = 0
for gold_label, inference_label in zip(
gold_labels, inference_labels):
if gold_label == inference_label:
if gold_label == 0:
right_0 += 1
elif gold_label == 1:
right_1 += 1
else:
right_2 += 1
elif inference_label == 0:
error_0 += 1
elif inference_label == 1:
error_1 += 1
else:
error_2 += 1
recall_0 = right_0 / (num_gold_0 + 1e-5)
recall_1 = right_1 / (num_gold_1 + 1e-5)
recall_2 = right_2 / (num_gold_2 + 1e-5)
precision_0 = right_0 / (error_0 + right_0 + 1e-5)
precision_1 = right_1 / (error_1 + right_1 + 1e-5)
precision_2 = right_2 / (error_2 + right_2 + 1e-5)
f10 = 2 * precision_0 * recall_0 / (precision_0 +
recall_0 + 1e-5)
f11 = 2 * precision_1 * recall_1 / (precision_1 +
recall_1 + 1e-5)
f12 = 2 * precision_2 * recall_2 / (precision_2 +
recall_2 + 1e-5)
output_dev_result_file = os.path.join(
args.output_dir, "dev_results.txt")
with open(output_dev_result_file, 'a',
encoding='utf-8') as f:
f.write('precision:' + str(precision_0) + ' ' +
str(precision_1) + ' ' + str(precision_2) +
'\n')
f.write('recall:' + str(recall_0) + ' ' +
str(recall_1) + ' ' + str(recall_2) + '\n')
f.write('f1:' + str(f10) + ' ' + str(f11) + ' ' +
str(f12) + '\n' + '\n')
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = accuracy(inference_logits, gold_labels)
# draw loss.
eval_F1.append(round(eval_accuracy, 4))
ax.append(step)
plt.plot(ax,
eval_F1,
label='eval_F1',
linewidth=1,
color='r',
marker='o',
markerfacecolor='blue',
markersize=2)
for a, b in zip(ax, eval_F1):
plt.text(a, b, b, ha='center', va='bottom', fontsize=8)
plt.savefig(args.output_dir + '/labeled.jpg')
result = {
'eval_loss': eval_loss,
'eval_F1': eval_accuracy,
'global_step': global_step,
'loss': train_loss
}
output_eval_file = os.path.join(args.output_dir,
"eval_results.txt")
with open(output_eval_file, "a") as writer:
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
writer.write('*' * 80)
writer.write('\n')
if eval_accuracy > best_acc and 'dev' in file:
print("=" * 80)
print("more accurate model arises, now best F1 = ",
eval_accuracy)
print("Saving Model......")
best_acc = eval_accuracy
# Save a trained model, only save the model it-self
model_to_save = model.module if hasattr(
model, 'module') else model
output_model_file = os.path.join(
args.output_dir, "pytorch_model.bin")
torch.save(model_to_save.state_dict(),
output_model_file)
print("=" * 80)
'''
if (step+1) / int(num_train_optimization_steps/10) > 9.5:
print("=" * 80)
print("End of training. Saving Model......")
# Save a trained model, only save the model it-self
model_to_save = model.module if hasattr(model, 'module') else model
output_model_file = os.path.join(args.output_dir, "pytorch_model_final_step.bin")
torch.save(model_to_save.state_dict(), output_model_file)
print("=" * 80)
'''
if args.do_test == 'yes':
start_time = time.time()
print(
'___________________now testing for best eval f1 model_________________________'
)
try:
del model
except:
pass
gc.collect()
args.do_train = 'no'
model = BertForSequenceClassification.from_pretrained(os.path.join(
args.output_dir, "pytorch_model.bin"),
args,
config=config)
model.half()
for layer in model.modules():
if isinstance(layer, torch.nn.modules.batchnorm._BatchNorm):
layer.float()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from "
"https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif args.n_gpu > 1:
model = torch.nn.DataParallel(model)
for file, flag in [('test.csv', 'test')]:
inference_labels = []
gold_labels = []
eval_examples = read_examples(os.path.join(args.data_dir, file),
is_training=False,
label_name=args.label_name)
eval_features = convert_examples_to_features(
eval_examples, tokenizer, args.max_seq_length, args.split_num,
False)
all_input_ids = torch.tensor(select_field(eval_features,
'input_ids'),
dtype=torch.long)
all_input_mask = torch.tensor(select_field(eval_features,
'input_mask'),
dtype=torch.long)
all_segment_ids = torch.tensor(select_field(
eval_features, 'segment_ids'),
dtype=torch.long)
all_label = torch.tensor([f.label for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(
input_ids=input_ids,
token_type_ids=segment_ids,
attention_mask=input_mask).detach().cpu().numpy()
# print('test_logits=', logits)
label_ids = label_ids.to('cpu').numpy()
inference_labels.append(logits)
gold_labels.append(label_ids)
gold_labels = np.concatenate(gold_labels, 0)
logits = np.concatenate(inference_labels, 0)
if flag == 'dev':
print(flag, accuracy(logits, gold_labels))
elif flag == 'test':
df = pd.read_csv(os.path.join(args.data_dir, file))
df['label_0'] = logits[:, 0]
df['label_1'] = logits[:, 1]
df['label_2'] = logits[:, 2]
df[['id', 'label_0', 'label_1',
'label_2']].to_csv(os.path.join(args.output_dir,
"sub.csv"),
index=False)
# df[['id', 'label_0', 'label_1']].to_csv(os.path.join(args.output_dir, "sub.csv"), index=False)
else:
raise ValueError('flag not in [dev, test]')
print('inference time usd = {}s'.format(time.time() - start_time))
'''
def train(args, train_dataset, model, tokenizer):
""" Train the model """
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter()
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_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}
]
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
warmup_steps = args.warmup_steps if args.warmup_steps >= 1 else int(t_total * args.warmup_steps)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=warmup_steps, t_total=t_total)
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(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
# train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
train_iterator = range(int(args.num_train_epochs))
set_seed(args) # Added here for reproductibility (even between python 2 and 3)
first_time = time.time()
best_result = 0.0
for idx_epoch in train_iterator:
# epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
epoch_iterator = train_dataloader
preds = None
out_label_ids = None
for step, batch in enumerate(epoch_iterator):
model.train()
batch = tuple(t.to(args.device) for t in batch)
input_ids, attention_mask, token_type_ids, labels = batch[0], batch[1], batch[2], batch[3]
inputs = {'input_ids': input_ids,
'attention_mask': attention_mask,
'token_type_ids': token_type_ids if args.model_type in ['bert', 'xlnet'] \
and not args.no_segment else None, # XLM and RoBERTa don't use segment_ids
'labels': labels}
outputs = model(**inputs)
loss = outputs[0] # model outputs are always tuple in pytorch-transformers (see doc)
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()
total_norm = torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
tr_loss += loss.item()
if preds is None:
preds = outputs[1].detach().cpu().numpy()
out_label_ids = inputs['labels'].detach().cpu().numpy()
else:
preds = np.append(preds, outputs[1].detach().cpu().numpy(), axis=0)
out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0)
if (step + 1) % args.gradient_accumulation_steps == 0:
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
# Log metrics
if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer, data_type="dev", prefix=global_step)
for key, value in results.items():
tb_writer.add_scalar('eval_{}'.format(key), value, global_step)
tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
# current loss
cur_loss = (tr_loss - logging_loss) / args.logging_steps
tb_writer.add_scalar('loss', cur_loss, global_step)
logging_loss = tr_loss
# print log
log_string = "Job_{}:".format(args.job_id)
log_string += " epoch={:<3d}".format(idx_epoch)
log_string += " step={:<8d}".format(global_step)
log_string += " batch={:<4d}".format(labels.shape[0])
log_string += " lr={:<10.7f}".format(scheduler.get_lr()[0])
log_string += " train_loss={:<8.5f}".format(cur_loss)
log_string += " |g|={:<10.7f}".format(total_norm)
# calculate accuracy
if args.output_mode == "classification":
preds = np.argmax(preds, axis=1)
elif args.output_mode == "regression":
preds = np.squeeze(preds)
result = compute_metrics(args.task_name, preds, out_label_ids)
for key in sorted(result.keys()):
log_string += " {}_{}={:<8.5f}".format("train", key, result[key])
log_string += " mins={:<9.2f}".format(float(time.time() - first_time) / 60)
logger.info(log_string)
preds = None
out_label_ids = None
if args.local_rank in [-1, 0] 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)
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
if args.max_steps > 0 and global_step > args.max_steps:
epoch_iterator.close()
break
if args.max_steps > 0 and global_step > args.max_steps:
train_iterator.close()
break
if args.local_rank == -1 and not args.evaluate_during_training and args.evaluate_after_training: # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer, data_type="dev", prefix=global_step)
metrics = result_for_sorting(args.task_name, results)
if metrics >= best_result:
best_result = metrics
# Save model checkpoint
output_dir = os.path.join(args.output_dir, 'best')
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)
if args.local_rank in [-1, 0]:
tb_writer.close()
return global_step, tr_loss / global_step
def train(args, train_dataset, model, tokenizer, context_model):
""" Train the model """
assert context_model
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_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}
]
# optimizer_grouped_parameters += [
# {'params': [p for n, p in context_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 context_model.named_parameters() 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)
scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
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(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
context_model.zero_grad()
train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
set_seed(args) # Added here for reproductibility (even between python 2 and 3)
cat_context_model = ContextModel(model, context_model)
for epoch_num in train_iterator:
model.train()
context_model.train()
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(args.device) for t in batch)
inputs = {'input_ids': batch[0],
'attention_mask': batch[1],
'token_type_ids': batch[2] if args.model_type in ['bert', 'xlnet'] else None, # XLM don't use segment_ids
'labels': batch[3],
}
outputs = cat_context_model.concat(**inputs)
# outputs = model(**inputs)
# context_outputs = context_model(**inputs)
loss = outputs # model outputs are always tuple in pytorch-transformers (see doc)
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()
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
else:
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
torch.nn.utils.clip_grad_norm_(context_model.parameters(), args.max_grad_norm)
tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0:
scheduler.step() # Update learning rate schedule
optimizer.step()
model.zero_grad()
context_model.zero_grad()
global_step += 1
if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
# Log metrics
logger.info("Step {0}, Loss: {1}".format(global_step, loss.item()))
if args.do_eval and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
eval_examples = read_cmsqa_examples(args.predict_file, is_training=True, num_choices=args.num_choices)
eval_features = convert_examples_to_features(
eval_examples, tokenizer, args.max_seq_length, True)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
all_input_ids = torch.tensor(select_field(eval_features, 'input_ids'), dtype=torch.long)
all_input_mask = torch.tensor(select_field(eval_features, 'input_mask'), dtype=torch.long)
all_segment_ids = torch.tensor(select_field(eval_features, 'segment_ids'), dtype=torch.long)
all_label = torch.tensor([f.label for f in eval_features], dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.per_gpu_eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
device = args.device
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = None
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss = cat_context_model.concat(input_ids, segment_ids, input_mask, label_ids)
logits = cat_context_model.concat(input_ids, segment_ids, input_mask)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
tmp_eval_accuracy = accuracy(logits, label_ids)
eval_loss += tmp_eval_loss.mean().item()
eval_accuracy += tmp_eval_accuracy
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = eval_accuracy / nb_eval_examples
result = {'eval_loss': eval_loss,
'eval_accuracy': eval_accuracy}
if args.do_train:
result['global_step'] = global_step
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "a") as writer:
writer.write('Epoch' + str(epoch_num) + '\n')
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
output_dir = os.path.join(args.output_dir, str(epoch_num))
os.makedirs(output_dir)
logger.info("Saving model checkpoint to %s", 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(output_dir)
return global_step, tr_loss / global_step
def main(rank, args):
#print('params: '," T_warm: ",T_warm," all_iteration: ",all_iteration," lr: ",lr)
#writer = SummaryWriter('./model_snapshot_error')
# cuda_list=range(cuda_num)
#cuda_list=[2,1,0]
#print('run')
random.seed(1)
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
torch.set_num_threads(1)
config = json.load(open(args.config_file, 'r', encoding="utf-8"))
model = Model_Hotpot2(args, config)
#model.network.to(cudaid)
param_optimizer = list(model.network.named_parameters())
param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]
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=config["training"]["learning_rate"])
scheduler = WarmupLinearSchedule(
optimizer,
warmup_steps=config["training"]["warmup_proportion"],
t_total=config["training"]["total_training_steps"])
#cudaid=dist.get_rank()
cudaid = rank
model.network.to(cudaid)
args.device = cudaid
#model.cuda(cudaid)
accumulation_steps = 16
#model = nn.DataParallel(model, device_ids=cuda_list)
print('rank: ', rank)
accum_batch_loss = 1
#train_file='train_ms_roberta_plain_pair_sample_shuffle.txt'
train_file = 'train_ms_transformer_xh_shuffle.txt'
#train_file='train_ms_roberta.txt'
#iterator=NewsIterator(batch_size=24, npratio=4)
#for epoch in range(0,100):
batch_t = 0
iteration = 0
#print('train...',cuda_list)
pre_batch_t = 177796
epoch = 0
model.train()
for epoch in range(0, 10):
#while True:
all_loss = 0
all_batch = 0
data_batch = utils.get_batch_dist(train_file, 8, rank)
for batch in data_batch:
#g,imp_index,label
#batch=imp_index , g , label
#print('batch: ',batch)
batch_t += 1
#batch[0] = batch[0].to(torch.device('cuda:0'))
# g=batch[0].to(torch.device('cuda:'+str(cudaid)))
g = batch[0].to(torch.device('cuda:' + str(cudaid)))
logit = model.network((g, batch[-1]), cudaid)
label = batch[-1].cuda(cudaid)
#print('logit: ',logit.shape)
#pos_node_idx = [i for i in range(batch[2].size(0)) if batch[1][i].item() != -1]
#print('????label:',batch[0].ndata['label'])
#pos_node_idx=[i for i in range(batch[0].ndata['label'].size(0)) if batch[0].ndata['label'][i].item()!=-1 ]
#print('pos_node_idx: ',pos_node_idx)
# logit=logit[pos_node_idx].reshape(-1,2)
logit = logit.reshape(-1, 2)
#print('logit: ',logit.shape)
loss = F.nll_loss(
F.log_softmax(
logit.view(-1, logit.size(-1)),
dim=-1,
dtype=torch.float32,
),
label.view(-1),
reduction='sum',
#ignore_index=self.padding_idx,
)
#sample_size=float(sample_size.sum())
#loss=loss.sum()/sample_size/math.log(2)
loss = loss / len(batch[1]) / math.log(2)
# sample_size=float(sample_size)
# loss=loss/sample_size/math.log(2)
#print(' batch_t: ',batch_t, ' epoch: ',epoch,' loss: ',float(loss))
accum_batch_loss += float(loss)
all_loss += float(loss)
all_batch += 1
loss = loss / accumulation_steps
loss.backward()
if (batch_t) % accumulation_steps == 0:
#print('candidate_id: ',candidate_id)
# total_norm=0
# for p in model.network.parameters():
# if p.grad==None:
# print('error: ',index,p.size(),p.grad)
# param_norm = p.grad.data.norm(2)
# total_norm += param_norm.item() ** 2
# total_norm = total_norm ** (1. / 2)
# total_clip_norm=0
# for p in model.network.parameters():
# if p.grad==None:
# print('error: ',index,p.size(),p.grad)
# param_norm = p.grad.data.norm(2)
# total_clip_norm += param_norm.item() ** 2
# total_clip_norm = total_clip_norm ** (1. / 2)
iteration += 1
#adjust_learning_rate(optimizer,iteration)
average_gradients(model)
torch.nn.utils.clip_grad_norm_(model.network.parameters(), 1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
#print(' batch_t: ',batch_t, ' iteration: ', iteration, ' epoch: ',epoch,' accum_batch_loss: ',accum_batch_loss/accumulation_steps, " total_norm: ", total_norm,' clip_norm: ',total_clip_norm)
if rank == 0:
print(' batch_t: ', batch_t, ' iteration: ', iteration,
' epoch: ', epoch, ' accum_batch_loss: ',
accum_batch_loss / accumulation_steps)
torch.save(model.network.state_dict(),
'./models/transformer_xh' + str(epoch) + '.pkl')
accum_batch_loss = 0
dist.barrier()
def run_train(args):
# --------- data
# processor = BertProcessor(vocab_path=config['bert_vocab_path'], do_lower_case=args.do_lower_case)
processor = BertProcessor()
label_list = processor.get_labels()
label2id = {label: i for i, label in enumerate(label_list)}
id2label = {i: label for i, label in enumerate(label_list)}
train_data = processor.get_train(config['data_dir'] /
f"{args.data_name}.train.pkl")
train_examples = processor.create_examples(
lines=train_data,
example_type='train',
cached_examples_file=config['data_dir'] /
f"cached_train_examples_{args.arch}")
train_features = processor.create_features(
examples=train_examples,
max_seq_len=args.train_max_seq_len,
cached_features_file=config['data_dir'] /
"cached_train_features_{}_{}".format(args.train_max_seq_len,
args.arch))
train_dataset = processor.create_dataset(train_features,
is_sorted=args.sorted)
if args.sorted:
train_sampler = SequentialSampler(train_dataset)
else:
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size)
valid_data = processor.get_dev(config['data_dir'] /
f"{args.data_name}.valid.pkl")
valid_examples = processor.create_examples(
lines=valid_data,
example_type='valid',
cached_examples_file=config['data_dir'] /
f"cached_valid_examples_{args.arch}")
valid_features = processor.create_features(
examples=valid_examples,
max_seq_len=args.eval_max_seq_len,
cached_features_file=config['data_dir'] /
"cached_valid_features_{}_{}".format(args.eval_max_seq_len, args.arch))
valid_dataset = processor.create_dataset(valid_features)
valid_sampler = SequentialSampler(valid_dataset)
valid_dataloader = DataLoader(valid_dataset,
sampler=valid_sampler,
batch_size=args.eval_batch_size)
# ------- model
logger.info("initializing model")
if args.resume_path:
args.resume_path = Path(args.resume_path)
model = BertForMultiLable.from_pretrained(args.resume_path,
num_labels=len(label_list))
else:
model = BertForMultiLable.from_pretrained('bert-large-uncased',
num_labels=len(label_list))
# model = BertForMultiLable.from_pretrained('bert-large-uncased', num_labels=len(label_list))
t_total = int(
len(train_dataloader) / args.gradient_accumulation_steps * args.epochs)
param_optimizer = list(model.named_parameters())
# no_decay = ['bias', 'LayerNorm.weight', 'gamma', 'beta']
no_decay = ['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':
args.weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
warmup_steps = int(t_total * args.warmup_proportion)
# optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon, betas=(0.5, 0.999))
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
lr_scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=warmup_steps,
t_total=t_total)
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)
# ---- callbacks
logger.info("initializing callbacks")
train_monitor = TrainingMonitor(file_dir=config['figure_dir'],
arch=args.arch)
model_checkpoint = ModelCheckpoint(checkpoint_dir=config['checkpoint_dir'],
mode=args.mode,
monitor=args.monitor,
arch=args.arch,
save_best_only=args.save_best)
# **************************** training model ***********************
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Num Epochs = %d", args.epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
trainer = Trainer(
n_gpu=args.n_gpu,
model=model,
epochs=args.epochs,
logger=logger,
criterion=BCEWithLogLoss(),
optimizer=optimizer,
lr_scheduler=lr_scheduler,
early_stopping=None,
training_monitor=train_monitor,
fp16=args.fp16,
resume_path=args.resume_path,
grad_clip=args.grad_clip,
model_checkpoint=model_checkpoint,
gradient_accumulation_steps=args.gradient_accumulation_steps,
batch_metrics=[AccuracyThresh(thresh=0.5),
F1Score()],
epoch_metrics=[
AUC(average='micro', task_type='binary'),
MultiLabelReport(id2label=id2label),
F1Score()
])
trainer.train(train_data=train_dataloader,
valid_data=valid_dataloader,
seed=args.seed)
def train(args, train_dataset, model, tokenizer):
""" Train the model """
if args.local_rank in [-1, 0]:
tb_writer = SummaryWriter()
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
train_sampler = RandomSampler(
train_dataset) if args.local_rank == -1 else DistributedSampler(
train_dataset)
train_dataloader = DataLoader(train_dataset,
sampler=train_sampler,
batch_size=args.train_batch_size)
if args.max_steps > 0:
t_total = args.max_steps
args.num_train_epochs = args.max_steps // (
len(train_dataloader) // args.gradient_accumulation_steps) + 1
else:
t_total = len(
train_dataloader
) // args.gradient_accumulation_steps * args.num_train_epochs
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_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
}]
optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate,
eps=args.adam_epsilon)
scheduler = WarmupLinearSchedule(optimizer,
warmup_steps=args.warmup_steps,
t_total=t_total)
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(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs)
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)
logger.info(" Total optimization steps = %d", t_total)
global_step = 0
tr_loss, logging_loss = 0.0, 0.0
model.zero_grad()
train_iterator = trange(int(args.num_train_epochs),
desc="Epoch",
disable=args.local_rank not in [-1, 0])
set_seed(
args) # Added here for reproducibility (even between python 2 and 3)
for _ in train_iterator:
epoch_iterator = tqdm(train_dataloader,
desc="Iteration",
disable=args.local_rank not in [-1, 0])
for step, batch in enumerate(epoch_iterator):
inputs, labels = mask_tokens(batch, tokenizer,
args) if args.mlm else (batch, batch)
inputs = inputs.to(args.device)
labels = labels.to(args.device)
model.train()
outputs = model(inputs,
masked_lm_labels=labels) if args.mlm else model(
inputs, labels=labels)
loss = outputs[
0] # model outputs are always tuple in pytorch-transformers (see doc)
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:
loss.backward()
tr_loss += loss.item()
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)
optimizer.step()
scheduler.step() # Update learning rate schedule
model.zero_grad()
global_step += 1
if args.local_rank in [
-1, 0
] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
# Log metrics
if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer)
for key, value in results.items():
tb_writer.add_scalar('eval_{}'.format(key), value,
global_step)
tb_writer.add_scalar('lr',
scheduler.get_lr()[0], global_step)
tb_writer.add_scalar('loss', (tr_loss - logging_loss) /
args.logging_steps, global_step)
logging_loss = tr_loss
if args.local_rank in [
-1, 0
] 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)
torch.save(args,
os.path.join(output_dir, 'training_args.bin'))
logger.info("Saving model checkpoint to %s", output_dir)
if args.max_steps > 0 and global_step > args.max_steps:
epoch_iterator.close()
break
if args.max_steps > 0 and global_step > args.max_steps:
train_iterator.close()
break
if args.local_rank in [-1, 0]:
tb_writer.close()
return global_step, tr_loss / global_step
print("{}: F1:{:.4f}, P:{:.4f}, R:{:.4f}".format(k, f1, p, r))
else:
os.makedirs(checkpoint_dir) if not os.path.exists(checkpoint_dir) else None
with open(os.path.join(checkpoint_dir, "args.json"), mode="w") as f:
json.dump(args, f)
train_loader, val_loader = prepare_loader(data["train"], args, mode="single")
optimization_steps = args["epoch"]
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"])
scheduler = WarmupLinearSchedule(optimizer, optimization_steps // 10, optimization_steps)
loss_fct = nn.CrossEntropyLoss()
model = model.to(device)
i2c = {v: k for k, v in constants.categories.items()}
best_f1 = 0
for epoch in range(1, args["epoch"] + 1):
model.train()
total_loss = 0
for n_step, batch in enumerate(train_loader):
ids, x, mask, y = batch
x, mask, y = x.to(device), mask.to(device), y.to(device)
optimizer.zero_grad()
model.zero_grad()
outputs = model(x, attention_mask=mask)
if args["model"] == "bert":
outputs = outputs[0]
