From 3f10b55ecf5e32e02b5e1c54fbcbedf8c20f2eb8 Mon Sep 17 00:00:00 2001 From: andrea Date: Thu, 11 Mar 2021 22:06:53 +0100 Subject: [PATCH] fix apex --fp16 --- code/run_squad.py | 38 +++++++++++++++----------------------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/code/run_squad.py b/code/run_squad.py index 9030104..3ab5308 100644 --- a/code/run_squad.py +++ b/code/run_squad.py @@ -19,6 +19,7 @@ import numpy as np import torch from torch.utils.data import DataLoader, TensorDataset +from apex import amp from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE, WEIGHTS_NAME, CONFIG_NAME from pytorch_pretrained_bert.modeling import BertForQuestionAnswering from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear @@ -919,8 +920,8 @@ def main(args): for lr in lrs: model = BertForQuestionAnswering.from_pretrained( args.model, cache_dir=PYTORCH_PRETRAINED_BERT_CACHE) - if args.fp16: - model.half() + + model.to(device) if n_gpu > 1: model = torch.nn.DataParallel(model) @@ -934,26 +935,17 @@ def main(args): if any(nd in n for nd in no_decay)], 'weight_decay': 0.0} ] - 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=lr, - 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 = BertAdam(optimizer_grouped_parameters, + optimizer = BertAdam(optimizer_grouped_parameters, lr=lr, warmup=args.warmup_proportion, t_total=num_train_optimization_steps) + + if args.fp16: + model, optimizer = amp.initialize( + model, optimizer, opt_level="O2", + keep_batchnorm_fp32=None, loss_scale="dynamic" + ) + tr_loss = 0 nb_tr_examples = 0 nb_tr_steps = 0 @@ -977,9 +969,10 @@ def main(args): tr_loss += loss.item() nb_tr_examples += input_ids.size(0) nb_tr_steps += 1 - if args.fp16: - optimizer.backward(loss) + with amp.scale_loss(loss, optimizer) as scaled_loss: + scaled_loss.backward() + loss = scaled_loss else: loss.backward() if (step + 1) % args.gradient_accumulation_steps == 0: @@ -1051,8 +1044,7 @@ def main(args): eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_example_index) eval_dataloader = DataLoader(eval_data, batch_size=args.eval_batch_size) model = BertForQuestionAnswering.from_pretrained(args.output_dir) - if args.fp16: - model.half() + model.to(device) na_prob_thresh = 1.0