def train(metric):
    """Training function."""

    logging.info('Now we are doing BERT classification training on %s!', ctx)
    optimizer_params = {'learning_rate': lr, 'epsilon': epsilon, 'wd': 0.01}
    try:
        trainer = gluon.Trainer(
            model.collect_params(),
            args.optimizer,
            optimizer_params,
            update_on_kvstore=False)
    except ValueError as e:
        print(e)
        warnings.warn(
            'AdamW optimizer is not found. Please consider upgrading to '
            'mxnet>=1.5.0. Now the original Adam optimizer is used instead.')
        trainer = gluon.Trainer(
            model.collect_params(),
            'adam',
            optimizer_params,
            update_on_kvstore=False)

    step_size = batch_size * accumulate if accumulate else batch_size
    num_train_steps = int(num_train_examples / step_size * args.epochs)
    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 model.collect_params().values() if p.grad_req != 'null'
    ]
    # Set grad_req if gradient accumulation is required
    if accumulate:
        for p in params:
            p.grad_req = 'add'

    for epoch_id in range(args.epochs):
        metric.reset()
        step_loss = 0
        tic = time.time()
        for batch_id, seqs in enumerate(train_data):
            # set grad to zero for gradient accumulation
            if accumulate:
                if batch_id % accumulate == 0:
                    model.collect_params().zero_grad()
                    step_num += 1
            else:
                step_num += 1
            # learning rate schedule
            if step_num < num_warmup_steps:
                new_lr = lr * step_num / num_warmup_steps
            else:
                offset = (step_num - num_warmup_steps) * lr / (
                    num_train_steps - num_warmup_steps)
                new_lr = lr - offset
            trainer.set_learning_rate(new_lr)
            # forward and backward
            with mx.autograd.record():
                input_ids, valid_length, type_ids, label = seqs
                out = model(
                    input_ids.as_in_context(ctx), type_ids.as_in_context(ctx),
                    valid_length.astype('float32').as_in_context(ctx))
                ls = loss_function(out, label.as_in_context(ctx)).mean()
            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_loss += ls.asscalar()
            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
        mx.nd.waitall()
        if task.task_name == 'MNLI':
            logging.info('On MNLI Matched: ')
            evaluate(dev_data_matched, metric)
            logging.info('On MNLI Mismatched: ')
            evaluate(dev_data_mismatched, metric)
        else:
            evaluate(dev_data, metric)

        # save params
        params_saved = os.path.join(output_dir,
                                    'model_bert_{0}_{1}.params'.format(task.task_name, epoch_id))
        model.save_parameters(params_saved)
        logging.info('params saved in : %s', params_saved)
        toc = time.time()
        logging.info('Time cost=%.2fs', toc - tic)
        tic = toc
def train ( metric ) : logging . info ( 'Now we are doing BERT classification training on %s!' , ctx ) optimizer_params = { 'learning_rate' : lr , 'epsilon' : epsilon , 'wd' : 0.01 } try : trainer = gluon . Trainer ( model . collect_params ( ) , args . optimizer , optimizer_params , update_on_kvstore = False ) except ValueError as e : print ( e ) warnings . warn ( 'AdamW optimizer is not found. Please consider upgrading to ' 'mxnet>=1.5.0. Now the original Adam optimizer is used instead.' ) trainer = gluon . Trainer ( model . collect_params ( ) , 'adam' , optimizer_params , update_on_kvstore = False ) step_size = batch_size * accumulate if accumulate else batch_size num_train_steps = int ( num_train_examples / step_size * args . epochs ) warmup_ratio = args . warmup_ratio num_warmup_steps = int ( num_train_steps * warmup_ratio ) step_num = 0 for _ , v in model . collect_params ( '.*beta|.*gamma|.*bias' ) . items ( ) : v . wd_mult = 0.0 params = [ p for p in model . collect_params ( ) . values ( ) if p . grad_req != 'null' ] if accumulate : for p in params : p . grad_req = 'add' for epoch_id in range ( args . epochs ) : metric . reset ( ) step_loss = 0 tic = time . time ( ) for batch_id , seqs in enumerate ( train_data ) : if accumulate : if batch_id % accumulate == 0 : model . collect_params ( ) . zero_grad ( ) step_num += 1 else : step_num += 1 if step_num < num_warmup_steps : new_lr = lr * step_num / num_warmup_steps else : offset = ( step_num - num_warmup_steps ) * lr / ( num_train_steps - num_warmup_steps ) new_lr = lr - offset trainer . set_learning_rate ( new_lr ) with mx . autograd . record ( ) : input_ids , valid_length , type_ids , label = seqs out = model ( input_ids . as_in_context ( ctx ) , type_ids . as_in_context ( ctx ) , valid_length . astype ( 'float32' ) . as_in_context ( ctx ) ) ls = loss_function ( out , label . as_in_context ( ctx ) ) . mean ( ) ls . backward ( ) 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_loss += ls . asscalar ( ) 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 mx . nd . waitall ( ) if task . task_name == 'MNLI' : logging . info ( 'On MNLI Matched: ' ) evaluate ( dev_data_matched , metric ) logging . info ( 'On MNLI Mismatched: ' ) evaluate ( dev_data_mismatched , metric ) else : evaluate ( dev_data , metric ) params_saved = os . path . join ( output_dir , 'model_bert_{0}_{1}.params' . format ( task . task_name , epoch_id ) ) model . save_parameters ( params_saved ) logging . info ( 'params saved in : %s' , params_saved ) toc = time . time ( ) logging . info ( 'Time cost=%.2fs' , toc - tic ) tic = toc
Training function.