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
        $kvstore->pull($name, out => $arg_list, priority  => -$index);
    }, $param_arrays, $grad_arrays);
}

func _update_params(
    ArrayRef[ArrayRef[AI::MXNet::NDArray]] $param_arrays,
    ArrayRef[ArrayRef[AI::MXNet::NDArray]] $grad_arrays,
    AI::MXNet::Updater                     $updater,
    Int                                    $num_device,
    Maybe[AI::MXNet::KVStore]              $kvstore=,