AI-MXNet

 view release on metacpan or  search on metacpan

lib/AI/MXNet/Module/Bucketing.pm  view on Meta::CPAN

package AI::MXNet::Module::Bucketing;
use Mouse;
use AI::MXNet::Function::Parameters;
use AI::MXNet::Base;

=encoding UTF-8

=head1 NAME

AI::MXNet::Module::Bucketing

=head1 SYNOPSIS

    my $buckets = [10, 20, 30, 40, 50, 60];
    my $start_label   = 1;
    my $invalid_label = 0;

    my ($train_sentences, $vocabulary) = tokenize_text(
        './data/ptb.train.txt', start_label => $start_label,
        invalid_label => $invalid_label
    );
    my ($validation_sentences) = tokenize_text(
        './data/ptb.test.txt', vocab => $vocabulary,
        start_label => $start_label, invalid_label => $invalid_label
    );
    my $data_train  = mx->rnn->BucketSentenceIter(
        $train_sentences, $batch_size, buckets => $buckets,
        invalid_label => $invalid_label
    );
    my $data_val    = mx->rnn->BucketSentenceIter(
        $validation_sentences, $batch_size, buckets => $buckets,
        invalid_label => $invalid_label
    );

    my $stack = mx->rnn->SequentialRNNCell();
    for my $i (0..$num_layers-1)
    {
        $stack->add(mx->rnn->LSTMCell(num_hidden => $num_hidden, prefix => "lstm_l${i}_"));
    }

    my $sym_gen = sub {
        my $seq_len = shift;
        my $data  = mx->sym->Variable('data');
        my $label = mx->sym->Variable('softmax_label');
        my $embed = mx->sym->Embedding(
            data => $data, input_dim => scalar(keys %$vocabulary),
            output_dim => $num_embed, name => 'embed'
        );
        $stack->reset;
        my ($outputs, $states) = $stack->unroll($seq_len, inputs => $embed, merge_outputs => 1);
        my $pred = mx->sym->Reshape($outputs, shape => [-1, $num_hidden]);
        $pred    = mx->sym->FullyConnected(data => $pred, num_hidden => scalar(keys %$vocabulary), name => 'pred');
        $label   = mx->sym->Reshape($label, shape => [-1]);
        $pred    = mx->sym->SoftmaxOutput(data => $pred, label => $label, name => 'softmax');
        return ($pred, ['data'], ['softmax_label']);
    };

    my $contexts;
    if(defined $gpus)
    {
        $contexts = [map { mx->gpu($_) } split(/,/, $gpus)];
    }
    else
    {
        $contexts = mx->cpu(0);
    }

    my $model = mx->mod->BucketingModule(
        sym_gen             => $sym_gen,
        default_bucket_key  => $data_train->default_bucket_key,
        context             => $contexts
    );

    $model->fit(
        $data_train,
        eval_data           => $data_val,
        eval_metric         => mx->metric->Perplexity($invalid_label),
        kvstore             => $kv_store,
        optimizer           => $optimizer,
        optimizer_params    => {
                                    learning_rate => $lr,
                                    momentum      => $mom,
                                    wd            => $wd,
                            },
        initializer         => mx->init->Xavier(factor_type => "in", magnitude => 2.34),
        num_epoch           => $num_epoch,
        batch_end_callback  => mx->callback->Speedometer($batch_size, $disp_batches),
        ($chkp_epoch ? (epoch_end_callback  => mx->rnn->do_rnn_checkpoint($stack, $chkp_prefix, $chkp_epoch)) : ())
    );

=head1 DESCRIPTION

    Implements the AI::MXNet::Module::Base API, and allows multiple
    symbols to be used depending on the `bucket_key` provided by each different
    mini-batch of data
=cut


=head2 new

    Parameters
    ----------
    $sym_gen : subref or any perl object that overloads &{} op
        A sub when called with a bucket key, returns a list with triple
        of ($symbol, $data_names, $label_names).
    $default_bucket_key : str or anything else
        The key for the default bucket.
    $logger : Logger
    $context : AI::MXNet::Context or array ref of AI::MXNet::Context objects
        Default is cpu(0)

lib/AI/MXNet/Module/Bucketing.pm  view on Meta::CPAN

method set_params(
    HashRef[AI::MXNet::NDArray] $arg_params,
    HashRef[AI::MXNet::NDArray] $aux_params,
    Bool                        $allow_missing=0,
    Bool                        $force_init=1,
    Bool                        $allow_extra=0
)
{
    if(not $allow_missing)
    {
        $self->init_params(
            arg_params    => $arg_params,    aux_params => $aux_params,
            allow_missing => $allow_missing, force_init => $force_init,
            allow_extra   => $allow_extra
        );
       return;
    }
    if($self->params_initialized and not $force_init)
    {
        AI::MXNet::Logging->warning(
            "Parameters already initialized and force_init=False. "
            ."set_params call ignored."
        );
        return;
    }
    $self->_curr_module->set_params(
        $arg_params, $aux_params,
        allow_missing => $allow_missing,
        force_init    => $force_init,
        allow_extra   => $allow_extra
    );
    # because we didn't update self._arg_params, they are dirty now.
    $self->_params_dirty(1);
    $self->params_initialized(1);
}

method init_params(
    AI::MXNet::Initializer             :$initializer=AI::MXNet::Initializer->Uniform(scale => 0.01),
    Maybe[HashRef[AI::MXNet::NDArray]] :$arg_params=,
    Maybe[HashRef[AI::MXNet::NDArray]] :$aux_params=,
    Bool                               :$allow_missing=0,
    Bool                               :$force_init=0,
    Bool                               :$allow_extra=0
)
{
    return if($self->params_initialized and not $force_init);
    assert($self->binded, 'call bind before initializing the parameters');
    $self->_curr_module->init_params(
        initializer   => $initializer,
        arg_params    => $arg_params,
        aux_params    => $aux_params,
        allow_missing => $allow_missing,
        force_init    => $force_init,
        allow_extra   => $allow_extra
    );
    $self->_params_dirty(0);
    $self->params_initialized(1);
}

method get_states(Bool $merge_multi_context=1)
{
    assert($self->binded and $self->params_initialized);
    $self->_curr_module->get_states($merge_multi_context);
}

method set_states(:$states=, :$value=)
{
    assert($self->binded and $self->params_initialized);
    $self->_curr_module->set_states(states => $states, value => $value);
}

=head2 bind

    Binding for a AI::MXNet::Module::Bucketing means setting up the buckets and bind the
    executor for the default bucket key. Executors corresponding to other keys are
    binded afterwards with switch_bucket.

    Parameters
    ----------
    :$data_shapes : ArrayRef[AI::MXNet::DataDesc|NameShape]
        This should correspond to the symbol for the default bucket.
    :$label_shapes= : Maybe[ArrayRef[AI::MXNet::DataDesc|NameShape]]
        This should correspond to the symbol for the default bucket.
    :$for_training : Bool
        Default is 1.
    :$inputs_need_grad : Bool
        Default is 0.
    :$force_rebind : Bool
        Default is 0.
    :$shared_module : AI::MXNet::Module::Bucketing
        Default is undef. This value is currently not used.
    :$grad_req : str, array ref of str, hash ref of str to str
        Requirement for gradient accumulation. Can be 'write', 'add', or 'null'
        (defaults to 'write').
        Can be specified globally (str) or for each argument (array ref, hash ref).
    :$bucket_key : str
        bucket key for binding. by default is to use the ->default_bucket_key
=cut

method bind(
    ArrayRef[AI::MXNet::DataDesc|NameShape]                   :$data_shapes,
    Maybe[ArrayRef[AI::MXNet::DataDesc|NameShape]]            :$label_shapes=,
    Bool                                                      :$for_training=1,
    Bool                                                      :$inputs_need_grad=0,
    Bool                                                      :$force_rebind=0,
    Maybe[AI::MXNet::BaseModule]                              :$shared_module=,
    Str|ArrayRef[Str]|HashRef[Str]                            :$grad_req='write',
    Maybe[Str]                                                :$bucket_key=
)
{
    # in case we already initialized params, keep it
    my ($arg_params, $aux_params);
    if($self->params_initialized)
    {
        ($arg_params, $aux_params) = $self->get_params;
    }

    # force rebinding is typically used when one want to switch from
    # training to prediction phase.
    $self->_reset_bind if $force_rebind;

    if($self->binded)
    {

lib/AI/MXNet/Module/Bucketing.pm  view on Meta::CPAN

        return;
    }

    $self->_curr_module->init_optimizer(
        kvstore           => $kvstore,
        optimizer         => $optimizer,
        optimizer_params  => $optimizer_params,
        force_init        => $force_init
    );
    for my $mod (values %{ $self->_buckets })
    {
        if($mod ne $self->_curr_module)
        {
            $mod->borrow_optimizer($self->_curr_module);
        }
    }
    $self->optimizer_initialized(1);
}

method prepare(AI::MXNet::DataBatch $data_batch)
{
    assert($self->binded and $self->params_initialized);
    ## perform bind if have not done so yet
    my $original_bucket_key = $self->_curr_bucket_key;
    $self->switch_bucket(
        bucket_key   => $data_batch->bucket_key,
        data_shapes  => $data_batch->provide_data,
        label_shapes => $data_batch->provide_label
    );
    # switch back
    $self->switch_bucket(bucket_key => $original_bucket_key);
}

method forward(
    AI::MXNet::DataBatch  $data_batch,
    Bool                 :$is_train=
)
{
    assert($self->binded and $self->params_initialized);
    $self->switch_bucket(
        bucket_key   => $data_batch->bucket_key,
        data_shapes  => $data_batch->provide_data,
        label_shapes => $data_batch->provide_label
    );
    $self->_curr_module->forward($data_batch, is_train => $is_train);
}

method backward(Maybe[ArrayRef[AI::MXNet::NDArray]|AI::MXNet::NDArray] $out_grads=)
{
    assert($self->binded and $self->params_initialized);
    $self->_curr_module->backward($out_grads);
}

method update()
{
    assert($self->binded and $self->params_initialized and $self->optimizer_initialized);
    $self->_params_dirty(1);
    $self->_curr_module->update;
}

method get_outputs(Bool $merge_multi_context=1)
{
    assert($self->binded and $self->params_initialized);
    return $self->_curr_module->get_outputs($merge_multi_context);
}

method get_input_grads(Bool $merge_multi_context=1)
{
    assert($self->binded and $self->params_initialized and $self->inputs_need_grad);
    return $self->_curr_module->get_input_grads($merge_multi_context);
}

method update_metric(
    AI::MXNet::EvalMetric $eval_metric,
    ArrayRef[AI::MXNet::NDArray] $labels
)
{
    assert($self->binded and $self->params_initialized);
    $self->_curr_module->update_metric($eval_metric, $labels);
}

method symbol()
{
    assert($self->binded);
    return $self->_curr_module->symbol;
}

method get_symbol()
{
    assert($self->binded);
    return $self->_buckets->{ $self->_default_bucket_key }->symbol;
}

method install_monitor(AI::MXNet::Monitor $mon)
{
    assert($self->binded);
    for my $mod (values %{ $self->_buckets })
    {
        $mod->install_monitor($mon);
    }
}

1;