| <!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"/> |
| <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: '' |
| }; |
| </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("/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://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS-MML_HTMLorMML"></script> --> |
| <!-- --> |
| <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></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="../../install/index.html">Install</a> |
| <a class="main-nav-link" href="../../tutorials/index.html">Tutorials</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="../../gluon/index.html">About</a></li> |
| <li><a class="main-nav-link" href="http://gluon.mxnet.io">Tutorials</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="../../api/python/index.html">Python</a></li> |
| <li><a class="main-nav-link" href="../../api/scala/index.html">Scala</a></li> |
| <li><a class="main-nav-link" href="../../api/r/index.html">R</a></li> |
| <li><a class="main-nav-link" href="../../api/julia/index.html">Julia</a></li> |
| <li><a class="main-nav-link" href="../../api/c++/index.html">C++</a></li> |
| <li><a class="main-nav-link" href="../../api/perl/index.html">Perl</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="../../faq/index.html">FAQ</a></li> |
| <li><a class="main-nav-link" href="../../architecture/index.html">Architecture</a></li> |
| <li><a class="main-nav-link" href="https://github.com/apache/incubator-mxnet/tree/1.0.0/example">Examples</a></li> |
| <li><a class="main-nav-link" href="../../model_zoo/index.html">Model Zoo</a></li> |
| </ul> |
| </span> |
| <a class="main-nav-link" href="https://github.com/dmlc/mxnet">Github</a> |
| <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="../../community/index.html">Community</a></li> |
| <li><a class="main-nav-link" href="../../community/contribute.html">Contribute</a></li> |
| <li><a class="main-nav-link" href="../../community/powered_by.html">Powered By</a></li> |
| </ul> |
| </span> |
| <a class="main-nav-link" href="http://discuss.mxnet.io">Discuss</a> |
| <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">Versions(1.0.0)<span class="caret"></span></a><ul id="package-dropdown-menu" class="dropdown-menu"><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/>1.1.0</a></li><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/versions/1.0.0/index.html>1.0.0</a></li><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/versions/0.12.1/index.html>0.12.1</a></li><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/versions/0.12.0/index.html>0.12.0</a></li><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/versions/0.11.0/index.html>0.11.0</a></li><li><a class="main-nav-link" href=https://mxnet.incubator.apache.org/versions/master/index.html>master</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="../../install/index.html">Install</a></li> |
| <li><a class="main-nav-link" href="../../tutorials/index.html">Tutorials</a></li> |
| <li class="dropdown-submenu"> |
| <a href="#" tabindex="-1">Community</a> |
| <ul class="dropdown-menu"> |
| <li><a href="../../community/index.html" tabindex="-1">Community</a></li> |
| <li><a href="../../community/contribute.html" tabindex="-1">Contribute</a></li> |
| <li><a href="../../community/powered_by.html" tabindex="-1">Powered By</a></li> |
| </ul> |
| </li> |
| <li class="dropdown-submenu"> |
| <a href="#" tabindex="-1">API</a> |
| <ul class="dropdown-menu"> |
| <li><a href="../../api/python/index.html" tabindex="-1">Python</a> |
| </li> |
| <li><a href="../../api/scala/index.html" tabindex="-1">Scala</a> |
| </li> |
| <li><a href="../../api/r/index.html" tabindex="-1">R</a> |
| </li> |
| <li><a href="../../api/julia/index.html" tabindex="-1">Julia</a> |
| </li> |
| <li><a href="../../api/c++/index.html" tabindex="-1">C++</a> |
| </li> |
| <li><a href="../../api/perl/index.html" tabindex="-1">Perl</a> |
| </li> |
| </ul> |
| </li> |
| <li class="dropdown-submenu"> |
| <a href="#" tabindex="-1">Docs</a> |
| <ul class="dropdown-menu"> |
| <li><a href="../../tutorials/index.html" tabindex="-1">Tutorials</a></li> |
| <li><a href="../../faq/index.html" tabindex="-1">FAQ</a></li> |
| <li><a href="../../architecture/index.html" tabindex="-1">Architecture</a></li> |
| <li><a href="https://github.com/apache/incubator-mxnet/tree/1.0.0/example" tabindex="-1">Examples</a></li> |
| <li><a href="../../model_zoo/index.html" tabindex="-1">Model Zoo</a></li> |
| </ul> |
| </li> |
| <li><a href="../../architecture/index.html">Architecture</a></li> |
| <li><a class="main-nav-link" href="https://github.com/dmlc/mxnet">Github</a></li> |
| <li id="dropdown-menu-position-anchor-version-mobile" class="dropdown-submenu" style="position: relative"><a href="#" tabindex="-1">Versions(1.0.0)</a><ul class="dropdown-menu"><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/>1.1.0</a></li><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/versions/1.0.0/index.html>1.0.0</a></li><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/versions/0.12.1/index.html>0.12.1</a></li><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/versions/0.12.0/index.html>0.12.0</a></li><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/versions/0.11.0/index.html>0.11.0</a></li><li><a tabindex="-1" href=https://mxnet.incubator.apache.org/versions/master/index.html>master</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="kn">import</span> <span class="n">absolute_import</span> |
| |
| <span class="kn">from</span> <span class="nn">array</span> <span class="kn">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="kn">import</span> <span class="n">NDArray</span> |
| <span class="kn">from</span> <span class="nn">.ndarray</span> <span class="kn">import</span> <span class="n">_ndarray_cls</span> |
| <span class="kn">from</span> <span class="nn">.base</span> <span class="kn">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="kn">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="kn">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="kn">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="bp">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="bp">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="fm">__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="bp">None</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_updater_func</span> <span class="o">=</span> <span class="bp">None</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">_str_updater_func</span> <span class="o">=</span> <span class="bp">None</span> |
| |
| <span class="k">def</span> <span class="fm">__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)=""></rowsparsendarray></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)=""></rowsparsendarray></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="bp">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="bp">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="bp">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="bp">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.</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="bp">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="bp">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="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></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="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></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="bp">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="bp">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="bp">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="bp">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="bp">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="bp">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, 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 type="text/javascript"> |
| $('body').ready(function () { |
| $('body').css('visibility', 'visible'); |
| }); |
| </script> |
| </body> |
| </html> |