Google Git
Sign in
apache / mxnet-site / 9bdb4340e6e81fd521d944fc8079eb0317ed3343 / . / versions / 1.1.0 / _modules / mxnet / kvstore.html
blob: 4951ca9366e06fe79d8ecd944168d526021f7b8b [file] [log] [blame]
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8"/>
<meta content="IE=edge" http-equiv="X-UA-Compatible"/>
<meta content="width=device-width, initial-scale=1" name="viewport"/>
<meta content="mxnet.kvstore" property="og:title">
<meta content="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/og-logo.png" property="og:image">
<meta content="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/og-logo.png" property="og:image:secure_url">
<meta content="mxnet.kvstore" property="og:description"/>
<title>mxnet.kvstore — mxnet documentation</title>
<link crossorigin="anonymous" href="https://maxcdn.bootstrapcdn.com/bootstrap/3.3.6/css/bootstrap.min.css" integrity="sha384-1q8mTJOASx8j1Au+a5WDVnPi2lkFfwwEAa8hDDdjZlpLegxhjVME1fgjWPGmkzs7" rel="stylesheet"/>
<link href="https://maxcdn.bootstrapcdn.com/font-awesome/4.5.0/css/font-awesome.min.css" rel="stylesheet"/>
<link href="../../_static/basic.css" rel="stylesheet" type="text/css">
<link href="../../_static/pygments.css" rel="stylesheet" type="text/css">
<link href="../../_static/mxnet.css" rel="stylesheet" type="text/css"/>
<script type="text/javascript">
var DOCUMENTATION_OPTIONS = {
URL_ROOT: '../../',
VERSION: '',
COLLAPSE_INDEX: false,
FILE_SUFFIX: '.html',
HAS_SOURCE: true,
SOURCELINK_SUFFIX: '.txt'
};
</script>
<script src="https://code.jquery.com/jquery-1.11.1.min.js" type="text/javascript"></script>
<script src="../../_static/underscore.js" type="text/javascript"></script>
<script src="../../_static/searchtools_custom.js" type="text/javascript"></script>
<script src="../../_static/doctools.js" type="text/javascript"></script>
<script src="../../_static/selectlang.js" type="text/javascript"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.1/MathJax.js?config=TeX-AMS-MML_HTMLorMML" type="text/javascript"></script>
<script type="text/javascript"> jQuery(function() { Search.loadIndex("/versions/1.1.0/searchindex.js"); Search.init();}); </script>
<script>
(function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
(i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new
Date();a=s.createElement(o),
m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
})(window,document,'script','https://www.google-analytics.com/analytics.js','ga');
ga('create', 'UA-96378503-1', 'auto');
ga('send', 'pageview');
</script>
<!-- -->
<!-- <script type="text/javascript" src="../../_static/jquery.js"></script> -->
<!-- -->
<!-- <script type="text/javascript" src="../../_static/underscore.js"></script> -->
<!-- -->
<!-- <script type="text/javascript" src="../../_static/doctools.js"></script> -->
<!-- -->
<!-- <script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.0/MathJax.js?config=TeX-AMS-MML_HTMLorMML"></script> -->
<!-- -->
<link href="../../genindex.html" rel="index" title="Index">
<link href="../../search.html" rel="search" title="Search"/>
<link href="../index.html" rel="up" title="Module code"/>
<link href="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/mxnet-icon.png" rel="icon" type="image/png"/>
</link></link></link></meta></meta></meta></head>
<body background="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/mxnet-background-compressed.jpeg" role="document">
<div class="content-block"><div class="navbar navbar-fixed-top">
<div class="container" id="navContainer">
<div class="innder" id="header-inner">
<h1 id="logo-wrap">
<a href="../../" id="logo"><img src="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/mxnet_logo.png"/></a>
</h1>
<nav class="nav-bar" id="main-nav">
<a class="main-nav-link" href="/versions/1.1.0/install/index.html">Install</a>
<span id="dropdown-menu-position-anchor">
<a aria-expanded="true" aria-haspopup="true" class="main-nav-link dropdown-toggle" data-toggle="dropdown" href="#" role="button">Gluon <span class="caret"></span></a>
<ul class="dropdown-menu navbar-menu" id="package-dropdown-menu">
<li><a class="main-nav-link" href="/versions/1.1.0/tutorials/gluon/gluon.html">About</a></li>
<li><a class="main-nav-link" href="https://www.d2l.ai/">Dive into Deep Learning</a></li>
<li><a class="main-nav-link" href="https://gluon-cv.mxnet.io">GluonCV Toolkit</a></li>
<li><a class="main-nav-link" href="https://gluon-nlp.mxnet.io/">GluonNLP Toolkit</a></li>
</ul>
</span>
<span id="dropdown-menu-position-anchor">
<a aria-expanded="true" aria-haspopup="true" class="main-nav-link dropdown-toggle" data-toggle="dropdown" href="#" role="button">API <span class="caret"></span></a>
<ul class="dropdown-menu navbar-menu" id="package-dropdown-menu">
<li><a class="main-nav-link" href="/versions/1.1.0/api/python/index.html">Python</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/c++/index.html">C++</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/julia/index.html">Julia</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/perl/index.html">Perl</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/r/index.html">R</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/scala/index.html">Scala</a></li>
</ul>
</span>
<span id="dropdown-menu-position-anchor-docs">
<a aria-expanded="true" aria-haspopup="true" class="main-nav-link dropdown-toggle" data-toggle="dropdown" href="#" role="button">Docs <span class="caret"></span></a>
<ul class="dropdown-menu navbar-menu" id="package-dropdown-menu-docs">
<li><a class="main-nav-link" href="/versions/1.1.0/faq/index.html">FAQ</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/tutorials/index.html">Tutorials</a>
<li><a class="main-nav-link" href="https://github.com/apache/incubator-mxnet/tree/v1.1.0/example">Examples</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/architecture/index.html">Architecture</a></li>
<li><a class="main-nav-link" href="https://cwiki.apache.org/confluence/display/MXNET/Apache+MXNet+Home">Developer Wiki</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/model_zoo/index.html">Model Zoo</a></li>
<li><a class="main-nav-link" href="https://github.com/onnx/onnx-mxnet">ONNX</a></li>
</li></ul>
</span>
<span id="dropdown-menu-position-anchor-community">
<a aria-expanded="true" aria-haspopup="true" class="main-nav-link dropdown-toggle" data-toggle="dropdown" href="#" role="button">Community <span class="caret"></span></a>
<ul class="dropdown-menu navbar-menu" id="package-dropdown-menu-community">
<li><a class="main-nav-link" href="http://discuss.mxnet.io">Forum</a></li>
<li><a class="main-nav-link" href="https://github.com/apache/incubator-mxnet/tree/v1.1.0">Github</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/community/contribute.html">Contribute</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/community/powered_by.html">Powered By</a></li>
</ul>
</span>
<span id="dropdown-menu-position-anchor-version" style="position: relative"><a href="#" class="main-nav-link dropdown-toggle" data-toggle="dropdown" role="button" aria-haspopup="true" aria-expanded="true">1.1.0<span class="caret"></span></a><ul id="package-dropdown-menu" class="dropdown-menu"><li><a href="/">master</a></li><li><a href=/versions/1.6/>1.6</a></li><li><a href=/versions/1.5.0/>1.5.0</a></li><li><a href=/versions/1.4.1/>1.4.1</a></li><li><a href=/versions/1.3.1/>1.3.1</a></li><li><a href=/versions/1.2.1/>1.2.1</a></li><li><a href=/versions/1.1.0/>1.1.0</a></li><li><a href=/versions/1.0.0/>1.0.0</a></li><li><a href=/versions/0.12.1/>0.12.1</a></li><li><a href=/versions/0.11.0/>0.11.0</a></li></ul></span></nav>
<script> function getRootPath(){ return "../../" } </script>
<div class="burgerIcon dropdown">
<a class="dropdown-toggle" data-toggle="dropdown" href="#" role="button">☰</a>
<ul class="dropdown-menu" id="burgerMenu">
<li><a href="/versions/1.1.0/install/index.html">Install</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/tutorials/index.html">Tutorials</a></li>
<li class="dropdown-submenu dropdown">
<a aria-expanded="true" aria-haspopup="true" class="dropdown-toggle burger-link" data-toggle="dropdown" href="#" tabindex="-1">Gluon</a>
<ul class="dropdown-menu navbar-menu" id="package-dropdown-menu">
<li><a class="main-nav-link" href="/versions/1.1.0/tutorials/gluon/gluon.html">About</a></li>
<li><a class="main-nav-link" href="http://gluon.mxnet.io">The Straight Dope (Tutorials)</a></li>
<li><a class="main-nav-link" href="https://gluon-cv.mxnet.io">GluonCV Toolkit</a></li>
<li><a class="main-nav-link" href="https://gluon-nlp.mxnet.io/">GluonNLP Toolkit</a></li>
</ul>
</li>
<li class="dropdown-submenu">
<a aria-expanded="true" aria-haspopup="true" class="dropdown-toggle burger-link" data-toggle="dropdown" href="#" tabindex="-1">API</a>
<ul class="dropdown-menu">
<li><a class="main-nav-link" href="/versions/1.1.0/api/python/index.html">Python</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/c++/index.html">C++</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/julia/index.html">Julia</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/perl/index.html">Perl</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/r/index.html">R</a></li>
<li><a class="main-nav-link" href="/versions/1.1.0/api/scala/index.html">Scala</a></li>
</ul>
</li>
<li class="dropdown-submenu">
<a aria-expanded="true" aria-haspopup="true" class="dropdown-toggle burger-link" data-toggle="dropdown" href="#" tabindex="-1">Docs</a>
<ul class="dropdown-menu">
<li><a href="/versions/1.1.0/faq/index.html" tabindex="-1">FAQ</a></li>
<li><a href="/versions/1.1.0/tutorials/index.html" tabindex="-1">Tutorials</a></li>
<li><a href="https://github.com/apache/incubator-mxnet/tree/v1.1.0/example" tabindex="-1">Examples</a></li>
<li><a href="/versions/1.1.0/architecture/index.html" tabindex="-1">Architecture</a></li>
<li><a href="https://cwiki.apache.org/confluence/display/MXNET/Apache+MXNet+Home" tabindex="-1">Developer Wiki</a></li>
<li><a href="/versions/1.1.0/model_zoo/index.html" tabindex="-1">Gluon Model Zoo</a></li>
<li><a href="https://github.com/onnx/onnx-mxnet" tabindex="-1">ONNX</a></li>
</ul>
</li>
<li class="dropdown-submenu dropdown">
<a aria-haspopup="true" class="dropdown-toggle burger-link" data-toggle="dropdown" href="#" role="button" tabindex="-1">Community</a>
<ul class="dropdown-menu">
<li><a href="http://discuss.mxnet.io" tabindex="-1">Forum</a></li>
<li><a href="https://github.com/apache/incubator-mxnet/tree/v1.1.0" tabindex="-1">Github</a></li>
<li><a href="/versions/1.1.0/community/contribute.html" tabindex="-1">Contribute</a></li>
<li><a href="/versions/1.1.0/community/powered_by.html" tabindex="-1">Powered By</a></li>
</ul>
</li>
<li id="dropdown-menu-position-anchor-version-mobile" class="dropdown-submenu" style="position: relative"><a href="#" tabindex="-1">1.1.0</a><ul class="dropdown-menu"><li><a tabindex="-1" href=/>master</a></li><li><a tabindex="-1" href=/versions/1.6/>1.6</a></li><li><a tabindex="-1" href=/versions/1.5.0/>1.5.0</a></li><li><a tabindex="-1" href=/versions/1.4.1/>1.4.1</a></li><li><a tabindex="-1" href=/versions/1.3.1/>1.3.1</a></li><li><a tabindex="-1" href=/versions/1.2.1/>1.2.1</a></li><li><a tabindex="-1" href=/versions/1.1.0/>1.1.0</a></li><li><a tabindex="-1" href=/versions/1.0.0/>1.0.0</a></li><li><a tabindex="-1" href=/versions/0.12.1/>0.12.1</a></li><li><a tabindex="-1" href=/versions/0.11.0/>0.11.0</a></li></ul></li></ul>
</div>
<div class="plusIcon dropdown">
<a class="dropdown-toggle" data-toggle="dropdown" href="#" role="button"><span aria-hidden="true" class="glyphicon glyphicon-plus"></span></a>
<ul class="dropdown-menu dropdown-menu-right" id="plusMenu"></ul>
</div>
<div id="search-input-wrap">
<form action="../../search.html" autocomplete="off" class="" method="get" role="search">
<div class="form-group inner-addon left-addon">
<i class="glyphicon glyphicon-search"></i>
<input class="form-control" name="q" placeholder="Search" type="text"/>
</div>
<input name="check_keywords" type="hidden" value="yes">
<input name="area" type="hidden" value="default"/>
</input></form>
<div id="search-preview"></div>
</div>
<div id="searchIcon">
<span aria-hidden="true" class="glyphicon glyphicon-search"></span>
</div>
<!-- <div id="lang-select-wrap"> -->
<!-- <label id="lang-select-label"> -->
<!-- <\!-- <i class="fa fa-globe"></i> -\-> -->
<!-- <span></span> -->
<!-- </label> -->
<!-- <select id="lang-select"> -->
<!-- <option value="en">Eng</option> -->
<!-- <option value="zh">中文</option> -->
<!-- </select> -->
<!-- </div> -->
<!-- <a id="mobile-nav-toggle">
<span class="mobile-nav-toggle-bar"></span>
<span class="mobile-nav-toggle-bar"></span>
<span class="mobile-nav-toggle-bar"></span>
</a> -->
</div>
</div>
</div>
<script type="text/javascript">
$('body').css('background', 'white');
</script>
<div class="container">
<div class="row">
<div aria-label="main navigation" class="sphinxsidebar leftsidebar" role="navigation">
<div class="sphinxsidebarwrapper">
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../api/python/index.html">Python Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/r/index.html">R Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/julia/index.html">Julia Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/c++/index.html">C++ Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/scala/index.html">Scala Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api/perl/index.html">Perl Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../faq/index.html">HowTo Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../architecture/index.html">System Documents</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../tutorials/index.html">Tutorials</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../community/index.html">Community</a></li>
</ul>
</div>
</div>
<div class="content">
<div class="page-tracker"></div>
<h1>Source code for mxnet.kvstore</h1><div class="highlight"><pre>
<span></span><span class="c1"># Licensed to the Apache Software Foundation (ASF) under one</span>
<span class="c1"># or more contributor license agreements. See the NOTICE file</span>
<span class="c1"># distributed with this work for additional information</span>
<span class="c1"># regarding copyright ownership. The ASF licenses this file</span>
<span class="c1"># to you under the Apache License, Version 2.0 (the</span>
<span class="c1"># "License"); you may not use this file except in compliance</span>
<span class="c1"># with the License. You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing,</span>
<span class="c1"># software distributed under the License is distributed on an</span>
<span class="c1"># "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY</span>
<span class="c1"># KIND, either express or implied. See the License for the</span>
<span class="c1"># specific language governing permissions and limitations</span>
<span class="c1"># under the License.</span>
<span class="c1"># coding: utf-8</span>
<span class="sd">""" Key value store interface of MXNet for parameter synchronization."""</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">absolute_import</span>
<span class="kn">from</span> <span class="nn">array</span> <span class="k">import</span> <span class="n">array</span>
<span class="kn">import</span> <span class="nn">ctypes</span>
<span class="kn">import</span> <span class="nn">pickle</span>
<span class="kn">from</span> <span class="nn">.ndarray</span> <span class="k">import</span> <span class="n">NDArray</span>
<span class="kn">from</span> <span class="nn">.ndarray</span> <span class="k">import</span> <span class="n">_ndarray_cls</span>
<span class="kn">from</span> <span class="nn">.base</span> <span class="k">import</span> <span class="n">_LIB</span><span class="p">,</span> <span class="n">c_str_array</span><span class="p">,</span> <span class="n">c_handle_array</span><span class="p">,</span> <span class="n">c_array</span><span class="p">,</span> <span class="n">c_array_buf</span><span class="p">,</span> <span class="n">c_str</span>
<span class="kn">from</span> <span class="nn">.base</span> <span class="k">import</span> <span class="n">check_call</span><span class="p">,</span> <span class="n">string_types</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">,</span> <span class="n">py_str</span>
<span class="kn">from</span> <span class="nn">.base</span> <span class="k">import</span> <span class="n">NDArrayHandle</span><span class="p">,</span> <span class="n">KVStoreHandle</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">optimizer</span> <span class="k">as</span> <span class="n">opt</span>
<span class="k">def</span> <span class="nf">_ctype_key_value</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">vals</span><span class="p">):</span>
<span class="sd">"""</span>
<span class="sd"> Returns ctype arrays for the key-value args, and the whether string keys are used.</span>
<span class="sd"> For internal use only.</span>
<span class="sd"> """</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="p">(</span><span class="nb">tuple</span><span class="p">,</span> <span class="nb">list</span><span class="p">)):</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">keys</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">vals</span><span class="p">))</span>
<span class="n">c_keys</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">c_vals</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">use_str_keys</span> <span class="o">=</span> <span class="kc">None</span>
<span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">val</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">vals</span><span class="p">):</span>
<span class="n">c_key_i</span><span class="p">,</span> <span class="n">c_val_i</span><span class="p">,</span> <span class="n">str_keys_i</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">val</span><span class="p">)</span>
<span class="n">c_keys</span> <span class="o">+=</span> <span class="n">c_key_i</span>
<span class="n">c_vals</span> <span class="o">+=</span> <span class="n">c_val_i</span>
<span class="n">use_str_keys</span> <span class="o">=</span> <span class="n">str_keys_i</span> <span class="k">if</span> <span class="n">use_str_keys</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">use_str_keys</span>
<span class="k">assert</span><span class="p">(</span><span class="n">use_str_keys</span> <span class="o">==</span> <span class="n">str_keys_i</span><span class="p">),</span> <span class="s2">"inconsistent types of keys detected."</span>
<span class="n">c_keys_arr</span> <span class="o">=</span> <span class="n">c_array</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_char_p</span><span class="p">,</span> <span class="n">c_keys</span><span class="p">)</span> <span class="k">if</span> <span class="n">use_str_keys</span> \
<span class="k">else</span> <span class="n">c_array</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">,</span> <span class="n">c_keys</span><span class="p">)</span>
<span class="n">c_vals_arr</span> <span class="o">=</span> <span class="n">c_array</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_void_p</span><span class="p">,</span> <span class="n">c_vals</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="n">c_keys_arr</span><span class="p">,</span> <span class="n">c_vals_arr</span><span class="p">,</span> <span class="n">use_str_keys</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="p">(</span><span class="nb">int</span><span class="p">,)</span> <span class="o">+</span> <span class="n">string_types</span><span class="p">)),</span> \
<span class="s2">"unexpected type for keys: "</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">keys</span><span class="p">))</span>
<span class="n">use_str_keys</span> <span class="o">=</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">keys</span><span class="p">,</span> <span class="n">string_types</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">vals</span><span class="p">,</span> <span class="n">NDArray</span><span class="p">):</span>
<span class="n">c_keys</span> <span class="o">=</span> <span class="n">c_str_array</span><span class="p">([</span><span class="n">keys</span><span class="p">])</span> <span class="k">if</span> <span class="n">use_str_keys</span> \
<span class="k">else</span> <span class="n">c_array_buf</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">,</span> <span class="n">array</span><span class="p">(</span><span class="s1">'i'</span><span class="p">,</span> <span class="p">[</span><span class="n">keys</span><span class="p">]))</span>
<span class="k">return</span> <span class="p">(</span><span class="n">c_keys</span><span class="p">,</span> <span class="n">c_handle_array</span><span class="p">([</span><span class="n">vals</span><span class="p">]),</span> <span class="n">use_str_keys</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">for</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">vals</span><span class="p">:</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">NDArray</span><span class="p">))</span>
<span class="n">c_keys</span> <span class="o">=</span> <span class="n">c_str_array</span><span class="p">([</span><span class="n">keys</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">vals</span><span class="p">))</span> <span class="k">if</span> <span class="n">use_str_keys</span> \
<span class="k">else</span> <span class="n">c_array_buf</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">,</span> <span class="n">array</span><span class="p">(</span><span class="s1">'i'</span><span class="p">,</span> <span class="p">[</span><span class="n">keys</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">vals</span><span class="p">)))</span>
<span class="k">return</span> <span class="p">(</span><span class="n">c_keys</span><span class="p">,</span> <span class="n">c_handle_array</span><span class="p">(</span><span class="n">vals</span><span class="p">),</span> <span class="n">use_str_keys</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_ctype_dict</span><span class="p">(</span><span class="n">param_dict</span><span class="p">):</span>
<span class="sd">"""</span>
<span class="sd"> Returns ctype arrays for keys and values(converted to strings) in a dictionary</span>
<span class="sd"> """</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">param_dict</span><span class="p">,</span> <span class="nb">dict</span><span class="p">)),</span> \
<span class="s2">"unexpected type for param_dict: "</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">param_dict</span><span class="p">))</span>
<span class="n">c_keys</span> <span class="o">=</span> <span class="n">c_array</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_char_p</span><span class="p">,</span> <span class="p">[</span><span class="n">c_str</span><span class="p">(</span><span class="n">k</span><span class="p">)</span> <span class="k">for</span> <span class="n">k</span> <span class="ow">in</span> <span class="n">param_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">()])</span>
<span class="n">c_vals</span> <span class="o">=</span> <span class="n">c_array</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">c_char_p</span><span class="p">,</span> <span class="p">[</span><span class="n">c_str</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">v</span><span class="p">))</span> <span class="k">for</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">param_dict</span><span class="o">.</span><span class="n">values</span><span class="p">()])</span>
<span class="k">return</span> <span class="p">(</span><span class="n">c_keys</span><span class="p">,</span> <span class="n">c_vals</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_updater_wrapper</span><span class="p">(</span><span class="n">updater</span><span class="p">):</span>
<span class="sd">"""A wrapper for the user-defined handle."""</span>
<span class="k">def</span> <span class="nf">updater_handle</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">lhs_handle</span><span class="p">,</span> <span class="n">rhs_handle</span><span class="p">,</span> <span class="n">_</span><span class="p">):</span>
<span class="sd">""" ctypes function """</span>
<span class="n">lhs</span> <span class="o">=</span> <span class="n">_ndarray_cls</span><span class="p">(</span><span class="n">NDArrayHandle</span><span class="p">(</span><span class="n">lhs_handle</span><span class="p">))</span>
<span class="n">rhs</span> <span class="o">=</span> <span class="n">_ndarray_cls</span><span class="p">(</span><span class="n">NDArrayHandle</span><span class="p">(</span><span class="n">rhs_handle</span><span class="p">))</span>
<span class="n">updater</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">lhs</span><span class="p">,</span> <span class="n">rhs</span><span class="p">)</span>
<span class="k">return</span> <span class="n">updater_handle</span>
<div class="viewcode-block" id="KVStore"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore">[docs]</a><span class="k">class</span> <span class="nc">KVStore</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
<span class="sd">"""A key-value store for synchronization of values, over multiple devices."""</span>
<span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">handle</span><span class="p">):</span>
<span class="sd">"""Initializes a new KVStore.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> handle : KVStoreHandle</span>
<span class="sd"> `KVStore` handle of C API.</span>
<span class="sd"> """</span>
<span class="k">assert</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">handle</span><span class="p">,</span> <span class="n">KVStoreHandle</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span> <span class="o">=</span> <span class="n">handle</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_updater</span> <span class="o">=</span> <span class="kc">None</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_updater_func</span> <span class="o">=</span> <span class="kc">None</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_str_updater_func</span> <span class="o">=</span> <span class="kc">None</span>
<span class="k">def</span> <span class="nf">__del__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreFree</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">))</span>
<div class="viewcode-block" id="KVStore.init"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.init">[docs]</a> <span class="k">def</span> <span class="nf">init</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
<span class="sd">""" Initializes a single or a sequence of key-value pairs into the store.</span>
<span class="sd"> For each key, one must `init` it before calling `push` or `pull`.</span>
<span class="sd"> When multiple workers invoke `init` for the same key, only</span>
<span class="sd"> the value supplied by worker with rank `0` is used. This function returns</span>
<span class="sd"> after data has been initialized successfully.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> key : str, int, or sequence of str or int</span>
<span class="sd"> The keys.</span>
<span class="sd"> value : NDArray, RowSparseNDArray or sequence of NDArray or RowSparseNDArray</span>
<span class="sd"> Values corresponding to the keys.</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> # init a single key-value pair</span>
<span class="sd"> >>> shape = (2,3)</span>
<span class="sd"> >>> kv = mx.kv.create('local')</span>
<span class="sd"> >>> kv.init('3', mx.nd.ones(shape)*2)</span>
<span class="sd"> >>> a = mx.nd.zeros(shape)</span>
<span class="sd"> >>> kv.pull('3', out=a)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 2. 2. 2.]</span>
<span class="sd"> [ 2. 2. 2.]]</span>
<span class="sd"> >>> # init a list of key-value pairs</span>
<span class="sd"> >>> keys = ['5', '7', '9']</span>
<span class="sd"> >>> kv.init(keys, [mx.nd.ones(shape)]*len(keys))</span>
<span class="sd"> >>> # init a row_sparse value</span>
<span class="sd"> >>> kv.init('4', mx.nd.ones(shape).tostype('row_sparse'))</span>
<span class="sd"> >>> b = mx.nd.sparse.zeros('row_sparse', shape)</span>
<span class="sd"> >>> kv.row_sparse_pull('4', row_ids=mx.nd.array([0, 1]), out=b)</span>
<span class="sd"> >>> print b</span>
<span class="sd"> <RowSparseNDArray 2x3 @cpu(0)></span>
<span class="sd"> """</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">use_str_keys</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
<span class="k">if</span> <span class="n">use_str_keys</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreInitEx</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">))</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreInit</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">))</span></div>
<div class="viewcode-block" id="KVStore.push"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.push">[docs]</a> <span class="k">def</span> <span class="nf">push</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">,</span> <span class="n">priority</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
<span class="sd">""" Pushes a single or a sequence of key-value pairs into the store.</span>
<span class="sd"> This function returns immediately after adding an operator to the engine.</span>
<span class="sd"> The actual operation is executed asynchronously. If there are consecutive</span>
<span class="sd"> pushes to the same key, there is no guarantee on the serialization of pushes.</span>
<span class="sd"> The execution of a push does not guarantee that all previous pushes are</span>
<span class="sd"> finished.</span>
<span class="sd"> There is no synchronization between workers.</span>
<span class="sd"> One can use ``_barrier()`` to sync all workers.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> key : str, int, or sequence of str or int</span>
<span class="sd"> Keys.</span>
<span class="sd"> value : NDArray, RowSparseNDArray, list of NDArray or RowSparseNDArray,</span>
<span class="sd"> or list of list of NDArray or RowSparseNDArray</span>
<span class="sd"> Values corresponding to the keys.</span>
<span class="sd"> priority : int, optional</span>
<span class="sd"> The priority of the push operation.</span>
<span class="sd"> Higher priority push operations are likely to be executed before</span>
<span class="sd"> other push actions.</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> # push a single key-value pair</span>
<span class="sd"> >>> kv.push('3', mx.nd.ones(shape)*8)</span>
<span class="sd"> >>> kv.pull('3', out=a) # pull out the value</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 8. 8. 8.]</span>
<span class="sd"> [ 8. 8. 8.]]</span>
<span class="sd"> >>> # aggregate the value and the push</span>
<span class="sd"> >>> gpus = [mx.gpu(i) for i in range(4)]</span>
<span class="sd"> >>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]</span>
<span class="sd"> >>> kv.push('3', b)</span>
<span class="sd"> >>> kv.pull('3', out=a)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 4. 4. 4.]</span>
<span class="sd"> [ 4. 4. 4.]]</span>
<span class="sd"> >>> # push a list of keys.</span>
<span class="sd"> >>> # single device</span>
<span class="sd"> >>> keys = ['4', '5', '6']</span>
<span class="sd"> >>> kv.push(keys, [mx.nd.ones(shape)]*len(keys))</span>
<span class="sd"> >>> b = [mx.nd.zeros(shape)]*len(keys)</span>
<span class="sd"> >>> kv.pull(keys, out=b)</span>
<span class="sd"> >>> print b[1].asnumpy()</span>
<span class="sd"> [[ 1. 1. 1.]</span>
<span class="sd"> [ 1. 1. 1.]]</span>
<span class="sd"> >>> # multiple devices:</span>
<span class="sd"> >>> keys = ['7', '8', '9']</span>
<span class="sd"> >>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)</span>
<span class="sd"> >>> kv.push(keys, b)</span>
<span class="sd"> >>> kv.pull(keys, out=b)</span>
<span class="sd"> >>> print b[1][1].asnumpy()</span>
<span class="sd"> [[ 4. 4. 4.]</span>
<span class="sd"> [ 4. 4. 4.]]</span>
<span class="sd"> >>> # push a row_sparse value</span>
<span class="sd"> >>> b = mx.nd.sparse.zeros('row_sparse', shape)</span>
<span class="sd"> >>> kv.init('10', mx.nd.sparse.zeros('row_sparse', shape))</span>
<span class="sd"> >>> kv.push('10', mx.nd.ones(shape).tostype('row_sparse'))</span>
<span class="sd"> >>> # pull out the value</span>
<span class="sd"> >>> kv.row_sparse_pull('10', row_ids=mx.nd.array([0, 1]), out=b)</span>
<span class="sd"> >>> print b</span>
<span class="sd"> <RowSparseNDArray 2x3 @cpu(0)></span>
<span class="sd"> """</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">use_str_keys</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">value</span><span class="p">)</span>
<span class="k">if</span> <span class="n">use_str_keys</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePushEx</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePush</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span></div>
<div class="viewcode-block" id="KVStore.pull"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.pull">[docs]</a> <span class="k">def</span> <span class="nf">pull</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">priority</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
<span class="sd">""" Pulls a single value or a sequence of values from the store.</span>
<span class="sd"> This function returns immediately after adding an operator to the engine.</span>
<span class="sd"> Subsequent attempts to read from the `out` variable will be blocked until the</span>
<span class="sd"> pull operation completes.</span>
<span class="sd"> `pull` is executed asynchronously after all previous `pull` calls and only</span>
<span class="sd"> the last `push` call for the same input key(s) are finished.</span>
<span class="sd"> The returned values are guaranteed to be the latest values in the store.</span>
<span class="sd"> For `RowSparseNDArray` values, this call is ignored,</span>
<span class="sd"> please use ``row_sparse_pull`` instead.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> key : str, int, or sequence of str or int</span>
<span class="sd"> Keys.</span>
<span class="sd"> out: NDArray or list of NDArray or list of list of NDArray</span>
<span class="sd"> Values corresponding to the keys.</span>
<span class="sd"> priority : int, optional</span>
<span class="sd"> The priority of the pull operation.</span>
<span class="sd"> Higher priority pull operations are likely to be executed before</span>
<span class="sd"> other pull actions.</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> # pull a single key-value pair</span>
<span class="sd"> >>> a = mx.nd.zeros(shape)</span>
<span class="sd"> >>> kv.pull('3', out=a)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 2. 2. 2.]</span>
<span class="sd"> [ 2. 2. 2.]]</span>
<span class="sd"> >>> # pull into multiple devices</span>
<span class="sd"> >>> b = [mx.nd.ones(shape, gpu) for gpu in gpus]</span>
<span class="sd"> >>> kv.pull('3', out=b)</span>
<span class="sd"> >>> print b[1].asnumpy()</span>
<span class="sd"> [[ 2. 2. 2.]</span>
<span class="sd"> [ 2. 2. 2.]]</span>
<span class="sd"> >>> # pull a list of key-value pairs.</span>
<span class="sd"> >>> # On single device</span>
<span class="sd"> >>> keys = ['5', '7', '9']</span>
<span class="sd"> >>> b = [mx.nd.zeros(shape)]*len(keys)</span>
<span class="sd"> >>> kv.pull(keys, out=b)</span>
<span class="sd"> >>> print b[1].asnumpy()</span>
<span class="sd"> [[ 2. 2. 2.]</span>
<span class="sd"> [ 2. 2. 2.]]</span>
<span class="sd"> >>> # On multiple devices</span>
<span class="sd"> >>> keys = ['6', '8', '10']</span>
<span class="sd"> >>> b = [[mx.nd.ones(shape, gpu) for gpu in gpus]] * len(keys)</span>
<span class="sd"> >>> kv.pull(keys, out=b)</span>
<span class="sd"> >>> print b[1][1].asnumpy()</span>
<span class="sd"> [[ 2. 2. 2.]</span>
<span class="sd"> [ 2. 2. 2.]]</span>
<span class="sd"> """</span>
<span class="k">assert</span><span class="p">(</span><span class="n">out</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">use_str_keys</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">out</span><span class="p">)</span>
<span class="k">if</span> <span class="n">use_str_keys</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePullEx</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePull</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span></div>
<div class="viewcode-block" id="KVStore.row_sparse_pull"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.row_sparse_pull">[docs]</a> <span class="k">def</span> <span class="nf">row_sparse_pull</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">,</span> <span class="n">out</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">priority</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">row_ids</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="sd">""" Pulls a single RowSparseNDArray value or a sequence of RowSparseNDArray values \</span>
<span class="sd"> from the store with specified row_ids. When there is only one row_id, KVStoreRowSparsePull \</span>
<span class="sd"> is invoked just once and the result is broadcast to all the rest of outputs.</span>
<span class="sd"> `row_sparse_pull` is executed asynchronously after all previous</span>
<span class="sd"> `pull`/`row_sparse_pull` calls and the last `push` call for the</span>
<span class="sd"> same input key(s) are finished.</span>
<span class="sd"> The returned values are guaranteed to be the latest values in the store.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> key : str, int, or sequence of str or int</span>
<span class="sd"> Keys.</span>
<span class="sd"> out: RowSparseNDArray or list of RowSparseNDArray or list of list of RowSparseNDArray</span>
<span class="sd"> Values corresponding to the keys. The stype is expected to be row_sparse</span>
<span class="sd"> priority : int, optional</span>
<span class="sd"> The priority of the pull operation.</span>
<span class="sd"> Higher priority pull operations are likely to be executed before</span>
<span class="sd"> other pull actions.</span>
<span class="sd"> row_ids : NDArray or list of NDArray</span>
<span class="sd"> The row_ids for which to pull for each value. Each row_id is an 1D NDArray \</span>
<span class="sd"> whose values don't have to be unique nor sorted.</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> shape = (3, 3)</span>
<span class="sd"> >>> kv.init('3', mx.nd.ones(shape).tostype('row_sparse'))</span>
<span class="sd"> >>> a = mx.nd.sparse.zeros('row_sparse', shape)</span>
<span class="sd"> >>> row_ids = mx.nd.array([0, 2], dtype='int64')</span>
<span class="sd"> >>> kv.row_sparse_pull('3', out=a, row_ids=row_ids)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 1. 1. 1.]</span>
<span class="sd"> [ 0. 0. 0.]</span>
<span class="sd"> [ 1. 1. 1.]]</span>
<span class="sd"> >>> duplicate_row_ids = mx.nd.array([2, 2], dtype='int64')</span>
<span class="sd"> >>> kv.row_sparse_pull('3', out=a, row_ids=duplicate_row_ids)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 0. 0. 0.]</span>
<span class="sd"> [ 0. 0. 0.]</span>
<span class="sd"> [ 1. 1. 1.]]</span>
<span class="sd"> >>> unsorted_row_ids = mx.nd.array([1, 0], dtype='int64')</span>
<span class="sd"> >>> kv.row_sparse_pull('3', out=a, row_ids=unsorted_row_ids)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 1. 1. 1.]</span>
<span class="sd"> [ 1. 1. 1.]</span>
<span class="sd"> [ 0. 0. 0.]]</span>
<span class="sd"> """</span>
<span class="k">assert</span><span class="p">(</span><span class="n">out</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="n">row_ids</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">row_ids</span><span class="p">,</span> <span class="n">NDArray</span><span class="p">):</span>
<span class="n">row_ids</span> <span class="o">=</span> <span class="p">[</span><span class="n">row_ids</span><span class="p">]</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">isinstance</span><span class="p">(</span><span class="n">row_ids</span><span class="p">,</span> <span class="nb">list</span><span class="p">)),</span> \
<span class="s2">"row_ids should be NDArray or list of NDArray"</span>
<span class="n">first_out</span> <span class="o">=</span> <span class="n">out</span>
<span class="c1"># whether row_ids are the same</span>
<span class="n">single_rowid</span> <span class="o">=</span> <span class="kc">False</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">row_ids</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">out</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
<span class="n">single_rowid</span> <span class="o">=</span> <span class="kc">True</span>
<span class="n">first_out</span> <span class="o">=</span> <span class="p">[</span><span class="n">out</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">use_str_keys</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">first_out</span><span class="p">)</span>
<span class="n">_</span><span class="p">,</span> <span class="n">crow_ids</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">_ctype_key_value</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">row_ids</span><span class="p">)</span>
<span class="k">assert</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">crow_ids</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">cvals</span><span class="p">)),</span> \
<span class="s2">"the number of row_ids doesn't match the number of values"</span>
<span class="k">if</span> <span class="n">use_str_keys</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePullRowSparseEx</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">crow_ids</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStorePullRowSparse</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">ckeys</span><span class="p">)),</span> <span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">,</span> <span class="n">crow_ids</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">(</span><span class="n">priority</span><span class="p">)))</span>
<span class="c1"># the result can be copied to other devices without invoking row_sparse_pull</span>
<span class="c1"># if the indices are the same</span>
<span class="k">if</span> <span class="n">single_rowid</span><span class="p">:</span>
<span class="k">for</span> <span class="n">out_i</span> <span class="ow">in</span> <span class="n">out</span><span class="p">[</span><span class="mi">1</span><span class="p">:]:</span>
<span class="n">out</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">copyto</span><span class="p">(</span><span class="n">out_i</span><span class="p">)</span></div>
<div class="viewcode-block" id="KVStore.set_gradient_compression"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.set_gradient_compression">[docs]</a> <span class="k">def</span> <span class="nf">set_gradient_compression</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">compression_params</span><span class="p">):</span>
<span class="sd">""" Specifies type of low-bit quantization for gradient compression \</span>
<span class="sd"> and additional arguments depending on the type of compression being used.</span>
<span class="sd"> 2bit Gradient Compression takes a positive float `threshold`.</span>
<span class="sd"> The technique works by thresholding values such that positive values in the</span>
<span class="sd"> gradient above threshold will be set to threshold. Negative values whose absolute</span>
<span class="sd"> values are higher than threshold, will be set to the negative of threshold.</span>
<span class="sd"> Values whose absolute values are less than threshold will be set to 0.</span>
<span class="sd"> By doing so, each value in the gradient is in one of three states. 2bits are</span>
<span class="sd"> used to represent these states, and every 16 float values in the original</span>
<span class="sd"> gradient can be represented using one float. This compressed representation</span>
<span class="sd"> can reduce communication costs. The difference between these thresholded values and</span>
<span class="sd"> original values is stored at the sender's end as residual and added to the</span>
<span class="sd"> gradient in the next iteration.</span>
<span class="sd"> When kvstore is 'local', gradient compression is used to reduce communication</span>
<span class="sd"> between multiple devices (gpus). Gradient is quantized on each GPU which</span>
<span class="sd"> computed the gradients, then sent to the GPU which merges the gradients. This</span>
<span class="sd"> receiving GPU dequantizes the gradients and merges them. Note that this</span>
<span class="sd"> increases memory usage on each GPU because of the residual array stored.</span>
<span class="sd"> When kvstore is 'dist', gradient compression is used to reduce communication</span>
<span class="sd"> from worker to sender. Gradient is quantized on each worker which</span>
<span class="sd"> computed the gradients, then sent to the server which dequantizes</span>
<span class="sd"> this data and merges the gradients from each worker. Note that this</span>
<span class="sd"> increases CPU memory usage on each worker because of the residual array stored.</span>
<span class="sd"> Only worker to server communication is compressed in this setting.</span>
<span class="sd"> If each machine has multiple GPUs, currently this GPU to GPU or GPU to CPU communication</span>
<span class="sd"> is not compressed. Server to worker communication (in the case of pull)</span>
<span class="sd"> is also not compressed.</span>
<span class="sd"> To use 2bit compression, we need to specify `type` as `2bit`.</span>
<span class="sd"> Only specifying `type` would use default value for the threshold.</span>
<span class="sd"> To completely specify the arguments for 2bit compression, we would need to pass</span>
<span class="sd"> a dictionary which includes `threshold` like:</span>
<span class="sd"> {'type': '2bit', 'threshold': 0.5}</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> compression_params : dict</span>
<span class="sd"> A dictionary specifying the type and parameters for gradient compression.</span>
<span class="sd"> The key `type` in this dictionary is a</span>
<span class="sd"> required string argument and specifies the type of gradient compression.</span>
<span class="sd"> Currently `type` can be only `2bit`</span>
<span class="sd"> Other keys in this dictionary are optional and specific to the type</span>
<span class="sd"> of gradient compression.</span>
<span class="sd"> """</span>
<span class="k">if</span> <span class="p">(</span><span class="s1">'device'</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">type</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="s1">'dist'</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">type</span><span class="p">):</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span> <span class="o">=</span> <span class="n">_ctype_dict</span><span class="p">(</span><span class="n">compression_params</span><span class="p">)</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreSetGradientCompression</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span>
<span class="n">mx_uint</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">compression_params</span><span class="p">)),</span>
<span class="n">ckeys</span><span class="p">,</span> <span class="n">cvals</span><span class="p">))</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s1">'Gradient compression is not supported for this type of kvstore'</span><span class="p">)</span></div>
<div class="viewcode-block" id="KVStore.set_optimizer"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.set_optimizer">[docs]</a> <span class="k">def</span> <span class="nf">set_optimizer</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">optimizer</span><span class="p">):</span>
<span class="sd">""" Registers an optimizer with the kvstore.</span>
<span class="sd"> When using a single machine, this function updates the local optimizer.</span>
<span class="sd"> If using multiple machines and this operation is invoked from a worker node,</span>
<span class="sd"> it will serialized the optimizer with pickle and send it to all servers.</span>
<span class="sd"> The function returns after all servers have been updated.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> optimizer : Optimizer</span>
<span class="sd"> The new optimizer for the store</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> kv = mx.kv.create()</span>
<span class="sd"> >>> shape = (2, 2)</span>
<span class="sd"> >>> weight = mx.nd.zeros(shape)</span>
<span class="sd"> >>> kv.init(3, weight)</span>
<span class="sd"> >>> # set the optimizer for kvstore as the default SGD optimizer</span>
<span class="sd"> >>> kv.set_optimizer(mx.optimizer.SGD())</span>
<span class="sd"> >>> grad = mx.nd.ones(shape)</span>
<span class="sd"> >>> kv.push(3, grad)</span>
<span class="sd"> >>> kv.pull(3, out = weight)</span>
<span class="sd"> >>> # weight is updated via gradient descent</span>
<span class="sd"> >>> weight.asnumpy()</span>
<span class="sd"> array([[-0.01, -0.01],</span>
<span class="sd"> [-0.01, -0.01]], dtype=float32)</span>
<span class="sd"> """</span>
<span class="n">is_worker</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">()</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreIsWorkerNode</span><span class="p">(</span><span class="n">ctypes</span><span class="o">.</span><span class="n">byref</span><span class="p">(</span><span class="n">is_worker</span><span class="p">)))</span>
<span class="c1"># pylint: disable=invalid-name</span>
<span class="k">if</span> <span class="s1">'dist'</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">type</span> <span class="ow">and</span> <span class="n">is_worker</span><span class="o">.</span><span class="n">value</span><span class="p">:</span>
<span class="c1"># send the optimizer to server</span>
<span class="k">try</span><span class="p">:</span>
<span class="c1"># use ASCII protocol 0, might be slower, but not a big ideal</span>
<span class="n">optim_str</span> <span class="o">=</span> <span class="n">py_str</span><span class="p">(</span><span class="n">pickle</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">optimizer</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
<span class="k">except</span><span class="p">:</span>
<span class="k">raise</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_send_command_to_servers</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">optim_str</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_set_updater</span><span class="p">(</span><span class="n">opt</span><span class="o">.</span><span class="n">get_updater</span><span class="p">(</span><span class="n">optimizer</span><span class="p">))</span></div>
<span class="nd">@property</span>
<span class="k">def</span> <span class="nf">type</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="sd">""" Returns the type of this kvstore.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> type : str</span>
<span class="sd"> the string type</span>
<span class="sd"> """</span>
<span class="n">kv_type</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_char_p</span><span class="p">()</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreGetType</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">byref</span><span class="p">(</span><span class="n">kv_type</span><span class="p">)))</span>
<span class="k">return</span> <span class="n">py_str</span><span class="p">(</span><span class="n">kv_type</span><span class="o">.</span><span class="n">value</span><span class="p">)</span>
<span class="nd">@property</span>
<span class="k">def</span> <span class="nf">rank</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="sd">""" Returns the rank of this worker node.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> rank : int</span>
<span class="sd"> The rank of this node, which is in range [0, num_workers())</span>
<span class="sd"> """</span>
<span class="n">rank</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">()</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreGetRank</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">byref</span><span class="p">(</span><span class="n">rank</span><span class="p">)))</span>
<span class="k">return</span> <span class="n">rank</span><span class="o">.</span><span class="n">value</span>
<span class="nd">@property</span>
<span class="k">def</span> <span class="nf">num_workers</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="sd">"""Returns the number of worker nodes.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> size :int</span>
<span class="sd"> The number of worker nodes.</span>
<span class="sd"> """</span>
<span class="n">size</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">()</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreGetGroupSize</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">byref</span><span class="p">(</span><span class="n">size</span><span class="p">)))</span>
<span class="k">return</span> <span class="n">size</span><span class="o">.</span><span class="n">value</span>
<div class="viewcode-block" id="KVStore.save_optimizer_states"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.save_optimizer_states">[docs]</a> <span class="k">def</span> <span class="nf">save_optimizer_states</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fname</span><span class="p">,</span> <span class="n">dump_optimizer</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
<span class="sd">"""Saves the optimizer (updater) state to a file. This is often used when checkpointing</span>
<span class="sd"> the model during training.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> fname : str</span>
<span class="sd"> Path to the output states file.</span>
<span class="sd"> dump_optimizer : bool, default False</span>
<span class="sd"> Whether to also save the optimizer itself. This would also save optimizer</span>
<span class="sd"> information such as learning rate and weight decay schedules.</span>
<span class="sd"> """</span>
<span class="k">assert</span> <span class="bp">self</span><span class="o">.</span><span class="n">_updater</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">"Cannot save states for distributed training"</span>
<span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">fname</span><span class="p">,</span> <span class="s1">'wb'</span><span class="p">)</span> <span class="k">as</span> <span class="n">fout</span><span class="p">:</span>
<span class="n">fout</span><span class="o">.</span><span class="n">write</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_updater</span><span class="o">.</span><span class="n">get_states</span><span class="p">(</span><span class="n">dump_optimizer</span><span class="p">))</span></div>
<div class="viewcode-block" id="KVStore.load_optimizer_states"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.KVStore.load_optimizer_states">[docs]</a> <span class="k">def</span> <span class="nf">load_optimizer_states</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">fname</span><span class="p">):</span>
<span class="sd">"""Loads the optimizer (updater) state from the file.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> fname : str</span>
<span class="sd"> Path to input states file.</span>
<span class="sd"> """</span>
<span class="k">assert</span> <span class="bp">self</span><span class="o">.</span><span class="n">_updater</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">"Cannot load states for distributed training"</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_updater</span><span class="o">.</span><span class="n">set_states</span><span class="p">(</span><span class="nb">open</span><span class="p">(</span><span class="n">fname</span><span class="p">,</span> <span class="s1">'rb'</span><span class="p">)</span><span class="o">.</span><span class="n">read</span><span class="p">())</span></div>
<span class="k">def</span> <span class="nf">_set_updater</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">updater</span><span class="p">):</span>
<span class="sd">"""Sets a push updater into the store.</span>
<span class="sd"> This function only changes the local store. When running on multiple machines one must</span>
<span class="sd"> use `set_optimizer`.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> updater : function</span>
<span class="sd"> The updater function.</span>
<span class="sd"> Examples</span>
<span class="sd"> --------</span>
<span class="sd"> >>> def update(key, input, stored):</span>
<span class="sd"> ... print "update on key: %d" % key</span>
<span class="sd"> ... stored += input * 2</span>
<span class="sd"> >>> kv._set_updater(update)</span>
<span class="sd"> >>> kv.pull('3', out=a)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 4. 4. 4.]</span>
<span class="sd"> [ 4. 4. 4.]]</span>
<span class="sd"> >>> kv.push('3', mx.nd.ones(shape))</span>
<span class="sd"> update on key: 3</span>
<span class="sd"> >>> kv.pull('3', out=a)</span>
<span class="sd"> >>> print a.asnumpy()</span>
<span class="sd"> [[ 6. 6. 6.]</span>
<span class="sd"> [ 6. 6. 6.]]</span>
<span class="sd"> """</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_updater</span> <span class="o">=</span> <span class="n">updater</span>
<span class="c1"># set updater with int keys</span>
<span class="n">_updater_proto</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">CFUNCTYPE</span><span class="p">(</span>
<span class="kc">None</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_int</span><span class="p">,</span> <span class="n">NDArrayHandle</span><span class="p">,</span> <span class="n">NDArrayHandle</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_void_p</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_updater_func</span> <span class="o">=</span> <span class="n">_updater_proto</span><span class="p">(</span><span class="n">_updater_wrapper</span><span class="p">(</span><span class="n">updater</span><span class="p">))</span>
<span class="c1"># set updater with str keys</span>
<span class="n">_str_updater_proto</span> <span class="o">=</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">CFUNCTYPE</span><span class="p">(</span>
<span class="kc">None</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_char_p</span><span class="p">,</span> <span class="n">NDArrayHandle</span><span class="p">,</span> <span class="n">NDArrayHandle</span><span class="p">,</span> <span class="n">ctypes</span><span class="o">.</span><span class="n">c_void_p</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_str_updater_func</span> <span class="o">=</span> <span class="n">_str_updater_proto</span><span class="p">(</span><span class="n">_updater_wrapper</span><span class="p">(</span><span class="n">updater</span><span class="p">))</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreSetUpdaterEx</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">_updater_func</span><span class="p">,</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_str_updater_func</span><span class="p">,</span> <span class="kc">None</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">_barrier</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="sd">"""Invokes global barrier among all worker nodes.</span>
<span class="sd"> For example, assume there are `n` machines. We would like machine `0` to first</span>
<span class="sd"> `init` the values and then have all the workers `pull` the initialized value.</span>
<span class="sd"> Before pulling, we can place invoke `_barrier()` to guarantee that the</span>
<span class="sd"> initialization is finished.</span>
<span class="sd"> """</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreBarrier</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">_send_command_to_servers</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">head</span><span class="p">,</span> <span class="n">body</span><span class="p">):</span>
<span class="sd">"""Sends a command to all server nodes.</span>
<span class="sd"> Sending command to a server node will cause that server node to invoke</span>
<span class="sd"> ``KVStoreServer.controller`` to execute the command.</span>
<span class="sd"> This function returns after the command has been executed on all server</span>
<span class="sd"> nodes.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> head : int</span>
<span class="sd"> the head of the command.</span>
<span class="sd"> body : str</span>
<span class="sd"> the body of the command.</span>
<span class="sd"> """</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreSendCommmandToServers</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">handle</span><span class="p">,</span> <span class="n">mx_uint</span><span class="p">(</span><span class="n">head</span><span class="p">),</span> <span class="n">c_str</span><span class="p">(</span><span class="n">body</span><span class="p">)))</span></div>
<div class="viewcode-block" id="create"><a class="viewcode-back" href="../../api/python/kvstore/kvstore.html#mxnet.kvstore.create">[docs]</a><span class="k">def</span> <span class="nf">create</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">'local'</span><span class="p">):</span>
<span class="sd">"""Creates a new KVStore.</span>
<span class="sd"> For single machine training, there are two commonly used types:</span>
<span class="sd"> ``local``: Copies all gradients to CPU memory and updates weights there.</span>
<span class="sd"> ``device``: Aggregates gradients and updates weights on GPUs. With this setting,</span>
<span class="sd"> the KVStore also attempts to use GPU peer-to-peer communication,</span>
<span class="sd"> potentially accelerating the communication.</span>
<span class="sd"> For distributed training, KVStore also supports a number of types:</span>
<span class="sd"> ``dist_sync``: Behaves similarly to ``local`` but with one major difference.</span>
<span class="sd"> With ``dist_sync``, batch-size now means the batch size used on each machine.</span>
<span class="sd"> So if there are ``n`` machines and we use batch size ``b``,</span>
<span class="sd"> then ``dist_sync`` behaves like ``local`` with batch size ``n * b``.</span>
<span class="sd"> ``dist_device_sync``: Identical to ``dist_sync`` with the difference similar</span>
<span class="sd"> to ``device`` vs ``local``.</span>
<span class="sd"> ``dist_async``: Performs asynchronous updates.</span>
<span class="sd"> The weights are updated whenever gradients are received from any machine.</span>
<span class="sd"> No two updates happen on the same weight at the same time. However, the order is not</span>
<span class="sd"> guaranteed.</span>
<span class="sd"> Parameters</span>
<span class="sd"> ----------</span>
<span class="sd"> name : {'local', 'device', 'nccl', 'dist_sync', 'dist_device_sync', 'dist_async'}</span>
<span class="sd"> The type of KVStore.</span>
<span class="sd"> Returns</span>
<span class="sd"> -------</span>
<span class="sd"> kv : KVStore</span>
<span class="sd"> The created KVStore.</span>
<span class="sd"> """</span>
<span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">name</span><span class="p">,</span> <span class="n">string_types</span><span class="p">):</span>
<span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="s1">'name must be a string'</span><span class="p">)</span>
<span class="n">handle</span> <span class="o">=</span> <span class="n">KVStoreHandle</span><span class="p">()</span>
<span class="n">check_call</span><span class="p">(</span><span class="n">_LIB</span><span class="o">.</span><span class="n">MXKVStoreCreate</span><span class="p">(</span><span class="n">c_str</span><span class="p">(</span><span class="n">name</span><span class="p">),</span>
<span class="n">ctypes</span><span class="o">.</span><span class="n">byref</span><span class="p">(</span><span class="n">handle</span><span class="p">)))</span>
<span class="k">return</span> <span class="n">KVStore</span><span class="p">(</span><span class="n">handle</span><span class="p">)</span></div>
</pre></div>
</div>
</div>
<div aria-label="main navigation" class="sphinxsidebar rightsidebar" role="navigation">
<div class="sphinxsidebarwrapper">
</div>
</div>
</div><div class="footer">
<div class="section-disclaimer">
<div class="container">
<div>
<img height="60" src="https://raw.githubusercontent.com/dmlc/web-data/master/mxnet/image/apache_incubator_logo.png"/>
<p>
Apache MXNet is an effort undergoing incubation at The Apache Software Foundation (ASF), <strong>sponsored by the <i>Apache Incubator</i></strong>. Incubation is required of all newly accepted projects until a further review indicates that the infrastructure, communications, and decision making process have stabilized in a manner consistent with other successful ASF projects. While incubation status is not necessarily a reflection of the completeness or stability of the code, it does indicate that the project has yet to be fully endorsed by the ASF.
</p>
<p>
"Copyright © 2017-2018, The Apache Software Foundation
Apache MXNet, MXNet, Apache, the Apache feather, and the Apache MXNet project logo are either registered trademarks or trademarks of the Apache Software Foundation."
</p>
</div>
</div>
</div>
</div> <!-- pagename != index -->
</div>
<script crossorigin="anonymous" integrity="sha384-0mSbJDEHialfmuBBQP6A4Qrprq5OVfW37PRR3j5ELqxss1yVqOtnepnHVP9aJ7xS" src="https://maxcdn.bootstrapcdn.com/bootstrap/3.3.6/js/bootstrap.min.js"></script>
<script src="../../_static/js/sidebar.js" type="text/javascript"></script>
<script src="../../_static/js/search.js" type="text/javascript"></script>
<script src="../../_static/js/navbar.js" type="text/javascript"></script>
<script src="../../_static/js/clipboard.min.js" type="text/javascript"></script>
<script src="../../_static/js/copycode.js" type="text/javascript"></script>
<script src="../../_static/js/page.js" type="text/javascript"></script>
<script src="../../_static/js/docversion.js" type="text/javascript"></script>
<script type="text/javascript">
$('body').ready(function () {
$('body').css('visibility', 'visible');
});
</script>
</body>
</html>