view release on metacpan or  search on metacpan

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

## TODO
## this class is here because of https://github.com/gfx/p5-Mouse/pull/67
## once 2.4.7 version of Mouse in Ubuntu for affected Perl version
## these accessors should be merged into main class

package AI::MXNet::Module::Private;
use Mouse;
has [qw/_param_names _fixed_param_names
        _aux_names _data_names _label_names _state_names
        _output_names _arg_params _aux_params
        _params_dirty _optimizer _kvstore
         _update_on_kvstore _updater _work_load_list
        _preload_opt_states _exec_group
        _data_shapes _label_shapes _context _grad_req/
] => (is => 'rw', init_arg => undef);

package AI::MXNet::Module;
use AI::MXNet::Base;
use AI::MXNet::Function::Parameters;
use List::Util qw(max);
use Data::Dumper ();
use Mouse;

func _create_kvstore(
    Maybe[Str|AI::MXNet::KVStore] $kvstore,
    Int                           $num_device,
    HashRef[AI::MXNet::NDArray]   $arg_params
)
{
    my $update_on_kvstore = 1;
    my $kv;
    if(defined $kvstore)
    {
        if(blessed $kvstore)
        {
            $kv = $kvstore;
        }
        else
        {
            # create kvstore using the string type
            if($num_device == 1 and $kvstore !~ /dist/)
            {
                # no need to use kv for single device and single machine
            }
            else
            {
                $kv = AI::MXNet::KVStore->create($kvstore);
                if($kvstore eq 'local')
                {
                    # automatically select a proper local
                    my $max_size = max(map { product(@{ $_->shape }) } values %{ $arg_params });
                    if($max_size > 1024 * 1024 * 16)
                    {
                        $update_on_kvstore = 0;
                    }
                }
            }
        }
    }

    $update_on_kvstore = 0 if not $kv;
    return ($kv, $update_on_kvstore);
}

func _initialize_kvstore(
    AI::MXNet::KVStore           :$kvstore,
    HashRef[AI::MXNet::NDArray]  :$arg_params,
    ArrayRef[Str]                :$param_names,
    Bool                         :$update_on_kvstore,
    ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] :$param_arrays
)
{
    enumerate(sub{
        my ($idx, $param_on_devs) = @_;
        my $name = $param_names->[$idx];
        $kvstore->init($name, $arg_params->{ $name });
        if($update_on_kvstore)
        {
            $kvstore->pull($name, out => $param_on_devs, priority => -$idx);
        }
    }, $param_arrays);
}

func _update_params_on_kvstore(
    ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] $param_arrays,
    ArrayRef[AI::MXNet::NDArray]|ArrayRef[ArrayRef[AI::MXNet::NDArray]] $grad_arrays,
    AI::MXNet::KVStore           $kvstore,
    ArrayRef[Str]                $param_names
)
{
    enumerate(sub{
        my ($index, $arg_list, $grad_list) = @_;
        if(ref $grad_list eq 'ARRAY' and not defined $grad_list->[0])
        {
            return;
        }
        my $name = $param_names->[$index];
        # push gradient, priority is negative index
        $kvstore->push($name, $grad_list, priority => -$index);
        # pull back the weights

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

    }
    if($self->binded)
    {
        $self->logger->warning('Already binded, ignoring bind()');
        return;
    }
    $self->for_training($for_training);
    $self->inputs_need_grad($inputs_need_grad);
    $self->binded(1);
    $self->_p->_grad_req($grad_req);

    if(not $for_training)
    {
        assert(not $inputs_need_grad);
    }
    ($data_shapes, $label_shapes) = $self->_parse_data_desc(
        $self->data_names, $self->label_names, $data_shapes, $label_shapes
    );
    $self->_p->_data_shapes($data_shapes);
    $self->_p->_label_shapes($label_shapes);
    my $shared_group;
    if($shared_module)
    {
        assert($shared_module->binded and $shared_module->params_initialized);
        $shared_group = $shared_module->_p->_exec_group;
    }

    $self->_p->_exec_group(
        AI::MXNet::DataParallelExecutorGroup->new(
            symbol            => $self->_symbol,
            contexts          => $self->_p->_context,
            workload          => $self->_p->_work_load_list,
            data_shapes       => $self->_p->_data_shapes,
            label_shapes      => $self->_p->_label_shapes,
            param_names       => $self->_p->_param_names,
            state_names       => $state_names,
            for_training      => $for_training,
            inputs_need_grad  => $inputs_need_grad,
            shared_group      => $shared_group,
            logger            => $self->logger,
            fixed_param_names => $self->_p->_fixed_param_names,
            grad_req          => $grad_req
        )
    );
    if($shared_module)
    {
        $self->params_initialized(1);
        $self->_p->_arg_params($shared_module->_p->_arg_params);
        $self->_p->_aux_params($shared_module->_p->_aux_params);
    }
    elsif($self->params_initialized)
    {
        # if the parameters are already initialized, we are re-binding
        # so automatically copy the already initialized params
        $self->_p->_exec_group->set_params($self->_p->_arg_params, $self->_p->_aux_params);
    }
    else
    {
        assert(not defined $self->_p->_arg_params and not $self->_p->_aux_params);
        my @param_arrays = (
            map { AI::MXNet::NDArray->zeros($_->[0]->shape, dtype => $_->[0]->dtype) }
            @{ $self->_p->_exec_group->_p->param_arrays }
        );
        my %arg_params;
        @arg_params{ @{ $self->_p->_param_names } } = @param_arrays;
        $self->_p->_arg_params(\%arg_params);
        my @aux_arrays = (
            map { AI::MXNet::NDArray->zeros($_->[0]->shape, dtype => $_->[0]->dtype) }
            @{ $self->_p->_exec_group->_p->aux_arrays }
        );
        my %aux_params;
        @aux_params{ @{ $self->_p->_aux_names } } = @aux_arrays;
        $self->_p->_aux_params(\%aux_params);
    }
    if($shared_module and $shared_module->optimizer_initialized)
    {
        $self->borrow_optimizer($shared_module)
    }
}

=head2 reshape

    Reshape the module for new input shapes.
    Parameters
    ----------
    :$data_shapes : ArrayRef[AI::MXNet::DataDesc]
        Typically is $data_iter->provide_data.
    :$label_shapes= : Maybe[ArrayRef[AI::MXNet::DataDesc]]
        Typically is $data_iter->provide_label.
=cut

method reshape(
    ArrayRef[AI::MXNet::DataDesc|NameShape]        :$data_shapes,
    Maybe[ArrayRef[AI::MXNet::DataDesc|NameShape]] :$label_shapes=
)
{
    assert($self->binded);
    ($data_shapes, $label_shapes) = $self->_parse_data_desc(
        $self->data_names, $self->label_names, $data_shapes, $label_shapes
    );
    $self->_p->_data_shapes($data_shapes);
    $self->_p->_label_shapes($label_shapes);
    $self->_p->_exec_group->reshape($self->_p->_data_shapes, $self->_p->_label_shapes);
}

method init_optimizer(
    Str|AI::MXNet::KVStore :$kvstore='local',
    Optimizer              :$optimizer='sgd',
    HashRef                :$optimizer_params={ learning_rate => 0.01 },
    Bool                   :$force_init=0
)
{
    assert($self->binded and $self->params_initialized);
    if($self->optimizer_initialized and not $force_init)
    {
        $self->logger->warning('optimizer already initialized, ignoring...');
        return;
    }
    if($self->_p->_params_dirty)
    {
        $self->_sync_params_from_devices;
    }

    my ($kvstore, $update_on_kvstore) = _create_kvstore(
        $kvstore,
        scalar(@{$self->_p->_context}),
        $self->_p->_arg_params
    );
    my $batch_size = $self->_p->_exec_group->_p->batch_size;
    if($kvstore and $kvstore->type =~ /dist/ and $kvstore->type =~ /_sync/)
    {
        $batch_size *= $kvstore->num_workers;
    }
    my $rescale_grad = 1/$batch_size;

    if(not blessed $optimizer)
    {
        my %idx2name;
        if($update_on_kvstore)
        {
            @idx2name{ 0..@{$self->_p->_exec_group->param_names}-1 } = @{$self->_p->_exec_group->param_names};
        }
        else
        {
            for my $k (0..@{$self->_p->_context}-1)
            {
                @idx2name{ map { $_ + $k } 0..@{$self->_p->_exec_group->param_names}-1 } = @{$self->_p->_exec_group->param_names};
            }
        }
        if(not exists $optimizer_params->{rescale_grad})
        {
            $optimizer_params->{rescale_grad} = $rescale_grad;
        }
        $optimizer = AI::MXNet::Optimizer->create(
            $optimizer,
            sym  => $self->symbol,
            param_idx2name => \%idx2name,
            %{ $optimizer_params }
        );
        if($optimizer->rescale_grad != $rescale_grad)
        {
            AI::MXNet::Logging->warning(
                "Optimizer created manually outside Module but rescale_grad "
                ."is not normalized to 1.0/batch_size/num_workers (%s vs. %s). "
                ."Is this intended?",
                $optimizer->rescale_grad, $rescale_grad
            );
        }
    }

    $self->_p->_optimizer($optimizer);
    $self->_p->_kvstore($kvstore);
    $self->_p->_update_on_kvstore($update_on_kvstore);
    $self->_p->_updater(undef);

    if($kvstore)
    {
        # copy initialized local parameters to kvstore
        _initialize_kvstore(
            kvstore           => $kvstore,
            param_arrays      => $self->_p->_exec_group->_p->param_arrays,
            arg_params        => $self->_p->_arg_params,
            param_names       => $self->_p->_param_names,
            update_on_kvstore => $update_on_kvstore
        );
    }
    if($update_on_kvstore)
    {
        $kvstore->set_optimizer($self->_p->_optimizer);
    }
    else
    {
        $self->_p->_updater(AI::MXNet::Optimizer->get_updater($optimizer));
    }
    $self->optimizer_initialized(1);

    if($self->_p->_preload_opt_states)
    {
        $self->load_optimizer_states($self->_p->_preload_opt_states);
        $self->_p->_preload_opt_states(undef);
    }
}

=head2 borrow_optimizer

    Borrow optimizer from a shared module. Used in bucketing, where exactly the same
    optimizer (esp. kvstore) is used.

    Parameters
    ----------
    shared_module : AI::MXNet::Module
=cut

method borrow_optimizer(AI::MXNet::Module $shared_module)
{
    assert($shared_module->optimizer_initialized);
    $self->_p->_optimizer($shared_module->_p->_optimizer);
    $self->_p->_kvstore($shared_module->_p->_kvstore);
    $self->_p->_update_on_kvstore($shared_module->_p->_update_on_kvstore);
    $self->_p->_updater($shared_module->_p->_updater);
    $self->optimizer_initialized(1);
}

method forward(
    AI::MXNet::DataBatch $data_batch,
    Maybe[Bool]         :$is_train=
)
{
    assert($self->binded and $self->params_initialized);

    my @curr_data_shapes = map { $_->shape } @{ $self->data_shapes };
    my @new_data_shapes  = map { $_->shape } @{ $data_batch->data };
    if(Data::Dumper->Dump(\@curr_data_shapes) ne Data::Dumper->Dump(\@new_data_shapes))
    {
        my $new_dshape;
        if($data_batch->can('provide_data') and $data_batch->provide_data)
        {
            $new_dshape = $data_batch->provide_data;
        }
        else
        {
            $new_dshape = [];
            zip(sub {
                my ($i, $shape) = @_;
                push @{ $new_dshape }, AI::MXNet::DataDesc->new(
                    $i->name, $shape, $i->dtype, $i->layout
                );
            }, $self->data_shapes, \@new_data_shapes);
        }
        my $new_lshape;
        if($data_batch->can('provide_label') and $data_batch->provide_label)
        {
            $new_lshape = $data_batch->provide_label;
        }
        elsif($data_batch->can('label') and $data_batch->label)
        {
            $new_lshape = [];
            zip(sub {
                my ($i, $j) = @_;
                push @{ $new_lshape }, AI::MXNet::DataDesc->new(
                    $i->name, $j->shape, $i->dtype, $i->layout
                );
            }, $self->label_shapes, $data_batch->label);
        }
        $self->reshape(data_shapes => $new_dshape, label_shapes => $new_lshape);
    }
    $self->_p->_exec_group->forward($data_batch, $is_train);
}

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

method update()
{
    assert($self->binded and $self->params_initialized and $self->optimizer_initialized);
    $self->_p->_params_dirty(1);
    if($self->_p->_update_on_kvstore)
    {
        _update_params_on_kvstore(
            $self->_p->_exec_group->_p->param_arrays,
            $self->_p->_exec_group->_p->grad_arrays,
            $self->_p->_kvstore,
            $self->_p->_exec_group->param_names
        );
    }
    else
    {
        _update_params(
            $self->_p->_exec_group->_p->param_arrays,
            $self->_p->_exec_group->_p->grad_arrays,
            $self->_p->_updater,
            scalar(@{ $self->_p->_context}),
            $self->_p->_kvstore,
            $self->_p->_exec_group->param_names
        );
    }
}

method get_outputs(Bool $merge_multi_context=1)
{
    assert($self->binded and $self->params_initialized);
    return $self->_p->_exec_group->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->_p->_exec_group->get_input_grads($merge_multi_context);
}

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

method set_states(:$states=, :$value=)
{