def train(args):
    """Training helper."""
    vocab, row, col, counts = get_train_data(args)
    model = GloVe(token_to_idx=vocab.token_to_idx, output_dim=args.emsize,
                  dropout=args.dropout, x_max=args.x_max, alpha=args.alpha,
                  weight_initializer=mx.init.Uniform(scale=1 / args.emsize))
    context = get_context(args)
    model.initialize(ctx=context)
    if not args.no_hybridize:
        model.hybridize(static_alloc=not args.no_static_alloc)

    optimizer_kwargs = dict(learning_rate=args.lr, eps=args.adagrad_eps)
    params = list(model.collect_params().values())
    try:
        trainer = mx.gluon.Trainer(params, 'groupadagrad', optimizer_kwargs)
    except ValueError:
        logging.warning('MXNet <= v1.3 does not contain '
                        'GroupAdaGrad support. Falling back to AdaGrad')
        trainer = mx.gluon.Trainer(params, 'adagrad', optimizer_kwargs)

    index_dtype = 'int32'
    if counts.shape[0] >= np.iinfo(np.int32).max:
        index_dtype = 'int64'
        logging.info('Co-occurrence matrix is large. '
                     'Using int64 to represent sample indices.')
    indices = mx.nd.arange(counts.shape[0], dtype=index_dtype)
    for epoch in range(args.epochs):
        # Logging variables
        log_wc = 0
        log_start_time = time.time()
        log_avg_loss = 0

        mx.nd.shuffle(indices, indices)  # inplace shuffle
        bs = args.batch_size
        num_batches = indices.shape[0] // bs
        for i in range(num_batches):
            batch_indices = indices[bs * i:bs * (i + 1)]
            ctx = context[i % len(context)]
            batch_row = row[batch_indices].as_in_context(ctx)
            batch_col = col[batch_indices].as_in_context(ctx)
            batch_counts = counts[batch_indices].as_in_context(ctx)
            with mx.autograd.record():
                loss = model(batch_row, batch_col, batch_counts)
                loss.backward()

            if len(context) == 1 or (i + 1) % len(context) == 0:
                trainer.step(batch_size=1)

            # Logging
            log_wc += loss.shape[0]
            log_avg_loss += loss.mean().as_in_context(context[0])
            if (i + 1) % args.log_interval == 0:
                # Forces waiting for computation by computing loss value
                log_avg_loss = log_avg_loss.asscalar() / args.log_interval
                wps = log_wc / (time.time() - log_start_time)
                logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, '
                             'throughput={:.2f}K wps, wc={:.2f}K'.format(
                                 epoch, i + 1, num_batches, log_avg_loss,
                                 wps / 1000, log_wc / 1000))
                log_dict = dict(
                    global_step=epoch * len(indices) + i * args.batch_size,
                    epoch=epoch, batch=i + 1, loss=log_avg_loss,
                    wps=wps / 1000)
                log(args, log_dict)

                log_start_time = time.time()
                log_avg_loss = 0
                log_wc = 0

            if args.eval_interval and (i + 1) % args.eval_interval == 0:
                with print_time('mx.nd.waitall()'):
                    mx.nd.waitall()
                with print_time('evaluate'):
                    evaluate(args, model, vocab, i + num_batches * epoch)

    # Evaluate
    with print_time('mx.nd.waitall()'):
        mx.nd.waitall()
    with print_time('evaluate'):
        evaluate(args, model, vocab, num_batches * args.epochs,
                 eval_analogy=not args.no_eval_analogy)

    # Save params
    with print_time('save parameters'):
        model.save_parameters(os.path.join(args.logdir, 'glove.params'))
def train ( args ) : vocab , row , col , counts = get_train_data ( args ) model = GloVe ( token_to_idx = vocab . token_to_idx , output_dim = args . emsize , dropout = args . dropout , x_max = args . x_max , alpha = args . alpha , weight_initializer = mx . init . Uniform ( scale = 1 / args . emsize ) ) context = get_context ( args ) model . initialize ( ctx = context ) if not args . no_hybridize : model . hybridize ( static_alloc = not args . no_static_alloc ) optimizer_kwargs = dict ( learning_rate = args . lr , eps = args . adagrad_eps ) params = list ( model . collect_params ( ) . values ( ) ) try : trainer = mx . gluon . Trainer ( params , 'groupadagrad' , optimizer_kwargs ) except ValueError : logging . warning ( 'MXNet <= v1.3 does not contain ' 'GroupAdaGrad support. Falling back to AdaGrad' ) trainer = mx . gluon . Trainer ( params , 'adagrad' , optimizer_kwargs ) index_dtype = 'int32' if counts . shape [ 0 ] >= np . iinfo ( np . int32 ) . max : index_dtype = 'int64' logging . info ( 'Co-occurrence matrix is large. ' 'Using int64 to represent sample indices.' ) indices = mx . nd . arange ( counts . shape [ 0 ] , dtype = index_dtype ) for epoch in range ( args . epochs ) : log_wc = 0 log_start_time = time . time ( ) log_avg_loss = 0 mx . nd . shuffle ( indices , indices ) bs = args . batch_size num_batches = indices . shape [ 0 ] // bs for i in range ( num_batches ) : batch_indices = indices [ bs * i : bs * ( i + 1 ) ] ctx = context [ i % len ( context ) ] batch_row = row [ batch_indices ] . as_in_context ( ctx ) batch_col = col [ batch_indices ] . as_in_context ( ctx ) batch_counts = counts [ batch_indices ] . as_in_context ( ctx ) with mx . autograd . record ( ) : loss = model ( batch_row , batch_col , batch_counts ) loss . backward ( ) if len ( context ) == 1 or ( i + 1 ) % len ( context ) == 0 : trainer . step ( batch_size = 1 ) log_wc += loss . shape [ 0 ] log_avg_loss += loss . mean ( ) . as_in_context ( context [ 0 ] ) if ( i + 1 ) % args . log_interval == 0 : log_avg_loss = log_avg_loss . asscalar ( ) / args . log_interval wps = log_wc / ( time . time ( ) - log_start_time ) logging . info ( '[Epoch {} Batch {}/{}] loss={:.4f}, ' 'throughput={:.2f}K wps, wc={:.2f}K' . format ( epoch , i + 1 , num_batches , log_avg_loss , wps / 1000 , log_wc / 1000 ) ) log_dict = dict ( global_step = epoch * len ( indices ) + i * args . batch_size , epoch = epoch , batch = i + 1 , loss = log_avg_loss , wps = wps / 1000 ) log ( args , log_dict ) log_start_time = time . time ( ) log_avg_loss = 0 log_wc = 0 if args . eval_interval and ( i + 1 ) % args . eval_interval == 0 : with print_time ( 'mx.nd.waitall()' ) : mx . nd . waitall ( ) with print_time ( 'evaluate' ) : evaluate ( args , model , vocab , i + num_batches * epoch ) with print_time ( 'mx.nd.waitall()' ) : mx . nd . waitall ( ) with print_time ( 'evaluate' ) : evaluate ( args , model , vocab , num_batches * args . epochs , eval_analogy = not args . no_eval_analogy ) with print_time ( 'save parameters' ) : model . save_parameters ( os . path . join ( args . logdir , 'glove.params' ) )
Training helper.