parser.add_argument('--fp16',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")

model.half()

# Prepare optimizerif 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': 0.01},{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]if args.fp16:try:from apex.optimizers import FP16_Optimizerfrom apex.optimizers import FusedAdamexcept 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)warmup_linear = WarmupLinearSchedule(warmup=args.warmup_proportion,t_total=num_train_optimization_steps)else:optimizer = BertAdam(optimizer_grouped_parameters,lr=args.learning_rate,warmup=args.warmup_proportion,t_total=num_train_optimization_steps)

 if args.fp16:optimizer.backward(loss)else:loss.backward()

if args.fp16:# modify learning rate with special warm up BERT uses# if args.fp16 is False, BertAdam is used that handles this automaticallylr_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_stepoptimizer.step()optimizer.zero_grad()