def build_optimizer(cfg: dict, net: gluon.HybridBlock):
lrs = build_lr_scheduler(cfg.pop('lr_scheduler', None))
cfg['optimizer_params']['lr_scheduler'] = lrs
net.backbone.collect_params().setattr('lr_mult',
cfg.pop('backbone_lr_mult', 1.0))
net.backbone.collect_params().setattr('wd_mult',
cfg.pop('backbone_wd_mult', 1.0))
if cfg.pop('no_wd', False):
net.collect_params('.*beta|.*gamma|.*bias').setattr('wd_mult', 0.0)
opt = cfg.pop('type', 'sgd')
optimizer_params = cfg.pop('optimizer_params', {})
if amp._amp_initialized:
cfg['update_on_kvstore'] = False
trainer = gluon.Trainer(net.collect_params(),
opt,
optimizer_params=optimizer_params,
**cfg)
if amp._amp_initialized:
amp.init_trainer(trainer)
return trainer
def train(metric):
"""Training function."""
if not only_inference:
logging.info('Now we are doing BERT classification training on %s!',
ctx)
all_model_params = model.collect_params()
optimizer_params = {'learning_rate': lr, 'epsilon': epsilon, 'wd': 0.01}
trainer = gluon.Trainer(all_model_params,
args.optimizer,
optimizer_params,
update_on_kvstore=False)
if args.dtype == 'float16':
amp.init_trainer(trainer)
epoch_number = args.epochs
step_size = batch_size * accumulate if accumulate else batch_size
num_train_steps = int(num_train_examples / step_size * args.epochs)
if args.training_steps:
num_train_steps = args.training_steps
epoch_number = 9999
logging.info('training steps=%d', num_train_steps)
warmup_ratio = args.warmup_ratio
num_warmup_steps = int(num_train_steps * warmup_ratio)
step_num = 0
# Do not apply weight decay on LayerNorm and bias terms
for _, v in model.collect_params('.*beta|.*gamma|.*bias').items():
v.wd_mult = 0.0
# Collect differentiable parameters
params = [p for p in all_model_params.values() if p.grad_req != 'null']
# Set grad_req if gradient accumulation is required
if accumulate and accumulate > 1:
for p in params:
p.grad_req = 'add'
# track best eval score
metric_history = []
best_metric = None
patience = args.early_stop
tic = time.time()
finish_flag = False
for epoch_id in range(epoch_number):
if args.early_stop and patience == 0:
logging.info('Early stopping at epoch %d', epoch_id)
break
if finish_flag:
break
if not only_inference:
metric.reset()
step_loss = 0
tic = time.time()
all_model_params.zero_grad()
for batch_id, seqs in enumerate(train_data):
# learning rate schedule
if step_num < num_warmup_steps:
new_lr = lr * step_num / num_warmup_steps
else:
non_warmup_steps = step_num - num_warmup_steps
offset = non_warmup_steps / (num_train_steps -
num_warmup_steps)
new_lr = lr - offset * lr
trainer.set_learning_rate(new_lr)
# forward and backward
with mx.autograd.record():
input_ids, segment_ids, valid_length, label = seqs
input_ids = input_ids.as_in_context(ctx)
valid_length = valid_length.as_in_context(ctx).astype(
'float32')
label = label.as_in_context(ctx)
if use_roberta:
out = model(input_ids, valid_length)
else:
out = model(input_ids, segment_ids.as_in_context(ctx),
valid_length)
ls = loss_function(out, label).mean()
if args.dtype == 'float16':
with amp.scale_loss(ls, trainer) as scaled_loss:
mx.autograd.backward(scaled_loss)
else:
ls.backward()
# update
if not accumulate or (batch_id + 1) % accumulate == 0:
trainer.allreduce_grads()
nlp.utils.clip_grad_global_norm(params, 1)
trainer.update(accumulate if accumulate else 1)
step_num += 1
if accumulate and accumulate > 1:
# set grad to zero for gradient accumulation
all_model_params.zero_grad()
step_loss += ls.asscalar()
if not do_regression:
label = label.reshape((-1))
metric.update([label], [out])
if (batch_id + 1) % (args.log_interval) == 0:
log_train(batch_id, len(train_data), metric, step_loss,
args.log_interval, epoch_id,
trainer.learning_rate)
step_loss = 0
if step_num >= num_train_steps:
logging.info('Finish training step: %d', step_num)
finish_flag = True
break
mx.nd.waitall()
# inference on dev data
for segment, dev_data in dev_data_list:
metric_nm, metric_val = evaluate(dev_data, metric, segment)
if best_metric is None or metric_val >= best_metric:
best_metric = metric_val
patience = args.early_stop
else:
if args.early_stop is not None:
patience -= 1
metric_history.append((epoch_id, metric_nm, metric_val))
if not only_inference:
# save params
ckpt_name = 'model_bert_{0}_{1}.params'.format(task_name, epoch_id)
params_saved = os.path.join(output_dir, ckpt_name)
nlp.utils.save_parameters(model, params_saved)
logging.info('params saved in: %s', params_saved)
toc = time.time()
logging.info('Time cost=%.2fs', toc - tic)
tic = toc
if not only_inference:
# we choose the best model based on metric[0],
# assuming higher score stands for better model quality
metric_history.sort(key=lambda x: x[2][0], reverse=True)
epoch_id, metric_nm, metric_val = metric_history[0]
ckpt_name = 'model_bert_{0}_{1}.params'.format(task_name, epoch_id)
params_saved = os.path.join(output_dir, ckpt_name)
nlp.utils.load_parameters(model, params_saved)
metric_str = 'Best model at epoch {}. Validation metrics:'.format(
epoch_id)
metric_str += ','.join([i + ':%.4f' for i in metric_nm])
logging.info(metric_str, *metric_val)
# inference on test data
for segment, test_data in test_data_list:
test(test_data, segment)