Example #1
0
 def cbow_batch(data):
     """Create a batch for CBOW training objective."""
     centers, word_context, word_context_mask = data
     assert len(centers.shape) == 2
     negatives_shape = (len(centers), args.negative)
     negatives, negatives_mask = remove_accidental_hits(
         negatives_sampler(negatives_shape), centers)
     center_negatives = mx.nd.concat(centers, negatives, dim=1)
     center_negatives_mask = mx.nd.concat(mx.nd.ones_like(centers),
                                          negatives_mask,
                                          dim=1)
     labels = mx.nd.concat(mx.nd.ones_like(centers),
                           mx.nd.zeros_like(negatives),
                           dim=1)
     if not args.ngram_buckets:
         return (word_context.as_in_context(context[0]),
                 word_context_mask.as_in_context(context[0]),
                 center_negatives.as_in_context(context[0]),
                 center_negatives_mask.as_in_context(context[0]),
                 labels.as_in_context(context[0]))
     else:
         unique, inverse_unique_indices = np.unique(word_context.asnumpy(),
                                                    return_inverse=True)
         inverse_unique_indices = mx.nd.array(inverse_unique_indices,
                                              ctx=context[0])
         subwords, subwords_mask = subword_lookup.get(unique.astype(int))
         return (word_context.as_in_context(context[0]),
                 word_context_mask.as_in_context(context[0]),
                 center_negatives.as_in_context(context[0]),
                 center_negatives_mask.as_in_context(context[0]),
                 labels.as_in_context(context[0]),
                 mx.nd.array(subwords, ctx=context[0]),
                 mx.nd.array(subwords_mask,
                             ctx=context[0]), inverse_unique_indices)
Example #2
0
def train(args):
    """Training helper."""
    if args.ngram_buckets:  # Fasttext model
        coded_dataset, negatives_sampler, vocab, subword_function, \
            idx_to_subwordidxs = get_train_data(args)
        embedding = nlp.model.train.FasttextEmbeddingModel(
            token_to_idx=vocab.token_to_idx,
            subword_function=subword_function,
            embedding_size=args.emsize,
            weight_initializer=mx.init.Uniform(scale=1 / args.emsize),
            sparse_grad=not args.no_sparse_grad,
        )
    else:
        coded_dataset, negatives_sampler, vocab = get_train_data(args)
        embedding = nlp.model.train.SimpleEmbeddingModel(
            token_to_idx=vocab.token_to_idx,
            embedding_size=args.emsize,
            weight_initializer=mx.init.Uniform(scale=1 / args.emsize),
            sparse_grad=not args.no_sparse_grad,
        )
    embedding_out = nlp.model.train.SimpleEmbeddingModel(
        token_to_idx=vocab.token_to_idx,
        embedding_size=args.emsize,
        weight_initializer=mx.init.Zero(),
        sparse_grad=not args.no_sparse_grad,
    )
    loss_function = mx.gluon.loss.SigmoidBinaryCrossEntropyLoss()

    context = get_context(args)
    embedding.initialize(ctx=context)
    embedding_out.initialize(ctx=context)
    if not args.no_hybridize:
        embedding.hybridize(static_alloc=not args.no_static_alloc)
        embedding_out.hybridize(static_alloc=not args.no_static_alloc)

    optimizer_kwargs = dict(learning_rate=args.lr)
    params = list(embedding.embedding.collect_params().values()) + \
        list(embedding_out.collect_params().values())
    trainer = mx.gluon.Trainer(params, args.optimizer, optimizer_kwargs)

    if args.ngram_buckets:  # Fasttext model
        optimizer_subwords_kwargs = dict(learning_rate=args.lr_subwords)
        params_subwords = list(
            embedding.subword_embedding.collect_params().values())
        trainer_subwords = mx.gluon.Trainer(params_subwords,
                                            args.optimizer_subwords,
                                            optimizer_subwords_kwargs)

    num_update = 0
    for epoch in range(args.epochs):
        random.shuffle(coded_dataset)
        context_sampler = nlp.data.ContextSampler(coded=coded_dataset,
                                                  batch_size=args.batch_size,
                                                  window=args.window)
        num_batches = len(context_sampler)

        # Logging variables
        log_wc = 0
        log_start_time = time.time()
        log_avg_loss = 0

        for i, batch in enumerate(context_sampler):
            progress = (epoch * num_batches + i) / (args.epochs * num_batches)
            (center, word_context, word_context_mask) = batch
            negatives_shape = (word_context.shape[0],
                               word_context.shape[1] * args.negative)
            negatives, negatives_mask = remove_accidental_hits(
                negatives_sampler(negatives_shape), word_context,
                word_context_mask)

            if args.ngram_buckets:  # Fasttext model
                if args.model.lower() == 'skipgram':
                    unique, inverse_unique_indices = np.unique(
                        center.asnumpy(), return_inverse=True)
                    unique = mx.nd.array(unique)
                    inverse_unique_indices = mx.nd.array(
                        inverse_unique_indices, ctx=context[0])
                    subwords, subwords_mask = \
                        indices_to_subwordindices_mask(unique, idx_to_subwordidxs)
                elif args.model.lower() == 'cbow':
                    unique, inverse_unique_indices = np.unique(
                        word_context.asnumpy(), return_inverse=True)
                    unique = mx.nd.array(unique)
                    inverse_unique_indices = mx.nd.array(
                        inverse_unique_indices, ctx=context[0])
                    subwords, subwords_mask = \
                        indices_to_subwordindices_mask(unique, idx_to_subwordidxs)
                else:
                    logging.error('Unsupported model %s.', args.model)
                    sys.exit(1)

            num_update += len(center)

            # To GPU
            center = center.as_in_context(context[0])
            if args.ngram_buckets:  # Fasttext model
                subwords = subwords.as_in_context(context[0])
                subwords_mask = subwords_mask.astype(np.float32).as_in_context(
                    context[0])
            word_context = word_context.as_in_context(context[0])
            word_context_mask = word_context_mask.as_in_context(context[0])
            negatives = negatives.as_in_context(context[0])
            negatives_mask = negatives_mask.as_in_context(context[0])

            with mx.autograd.record():
                # Combine subword level embeddings with word embeddings
                if args.model.lower() == 'skipgram':
                    if args.ngram_buckets:
                        emb_in = embedding(center, subwords,
                                           subwordsmask=subwords_mask,
                                           words_to_unique_subwords_indices=
                                           inverse_unique_indices)
                    else:
                        emb_in = embedding(center)

                    with mx.autograd.pause():
                        word_context_negatives = mx.nd.concat(
                            word_context, negatives, dim=1)
                        word_context_negatives_mask = mx.nd.concat(
                            word_context_mask, negatives_mask, dim=1)

                    emb_out = embedding_out(word_context_negatives,
                                            word_context_negatives_mask)

                    # Compute loss
                    pred = mx.nd.batch_dot(emb_in, emb_out.swapaxes(1, 2))
                    pred = pred.squeeze() * word_context_negatives_mask
                    label = mx.nd.concat(word_context_mask,
                                         mx.nd.zeros_like(negatives), dim=1)

                elif args.model.lower() == 'cbow':
                    word_context = word_context.reshape((-3, 1))
                    word_context_mask = word_context_mask.reshape((-3, 1))
                    if args.ngram_buckets:
                        emb_in = embedding(word_context, subwords,
                                           word_context_mask, subwords_mask,
                                           inverse_unique_indices)
                    else:
                        emb_in = embedding(word_context, word_context_mask)

                    with mx.autograd.pause():
                        center = center.tile(args.window * 2).reshape((-1, 1))
                        negatives = negatives.reshape((-1, args.negative))

                        center_negatives = mx.nd.concat(
                            center, negatives, dim=1)
                        center_negatives_mask = mx.nd.concat(
                            mx.nd.ones_like(center), negatives_mask, dim=1)

                    emb_out = embedding_out(center_negatives,
                                            center_negatives_mask)

                    # Compute loss
                    pred = mx.nd.batch_dot(emb_in, emb_out.swapaxes(1, 2))
                    pred = pred.squeeze() * word_context_mask
                    label = mx.nd.concat(
                        mx.nd.ones_like(word_context),
                        mx.nd.zeros_like(negatives), dim=1)

                loss = loss_function(pred, label)

            loss.backward()

            if args.optimizer.lower() != 'adagrad':
                trainer.set_learning_rate(
                    max(0.0001, args.lr * (1 - progress)))

            if (args.optimizer_subwords.lower() != 'adagrad'
                    and args.ngram_buckets):
                trainer_subwords.set_learning_rate(
                    max(0.0001, args.lr_subwords * (1 - progress)))

            trainer.step(batch_size=1)
            if args.ngram_buckets:
                trainer_subwords.step(batch_size=1)

            # Logging
            log_wc += loss.shape[0]
            log_avg_loss += loss.mean()
            if (i + 1) % args.log_interval == 0:
                wps = log_wc / (time.time() - log_start_time)
                # Forces waiting for computation by computing loss value
                log_avg_loss = log_avg_loss.asscalar() / args.log_interval
                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_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, embedding, vocab, num_update)

    # Evaluate
    with print_time('mx.nd.waitall()'):
        mx.nd.waitall()
    with print_time('evaluate'):
        evaluate(args, embedding, vocab, num_update,
                 eval_analogy=not args.no_eval_analogy)

    # Save params
    with print_time('save parameters'):
        save_params(args, embedding, embedding_out)