| { |
| "cells": [ |
| { |
| "cell_type": "markdown", |
| "id": "9a2dea4c", |
| "metadata": {}, |
| "source": [ |
| "<!--- Licensed to the Apache Software Foundation (ASF) under one -->\n", |
| "<!--- or more contributor license agreements. See the NOTICE file -->\n", |
| "<!--- distributed with this work for additional information -->\n", |
| "<!--- regarding copyright ownership. The ASF licenses this file -->\n", |
| "<!--- to you under the Apache License, Version 2.0 (the -->\n", |
| "<!--- \"License\"); you may not use this file except in compliance -->\n", |
| "<!--- with the License. You may obtain a copy of the License at -->\n", |
| "\n", |
| "<!--- http://www.apache.org/licenses/LICENSE-2.0 -->\n", |
| "\n", |
| "<!--- Unless required by applicable law or agreed to in writing, -->\n", |
| "<!--- software distributed under the License is distributed on an -->\n", |
| "<!--- \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->\n", |
| "<!--- KIND, either express or implied. See the License for the -->\n", |
| "<!--- specific language governing permissions and limitations -->\n", |
| "<!--- under the License. -->\n", |
| "\n", |
| "# Distributed Key-Value Store\n", |
| "\n", |
| "KVStore is a place for data sharing. Think of it as a single object shared\n", |
| "across different devices (GPUs and computers), where each device can push data in\n", |
| "and pull data out.\n", |
| "\n", |
| "## Initialization\n", |
| "\n", |
| "Let's consider a simple example: initializing\n", |
| "a (`int`, `NDArray`) pair into the store, and then pulling the value out:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "50bd5a20", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "import mxnet as mx\n", |
| "\n", |
| "kv = mx.kv.create('local') # create a local kv store.\n", |
| "shape = (2,3)\n", |
| "kv.init(3, mx.nd.ones(shape)*2)\n", |
| "a = mx.nd.zeros(shape)\n", |
| "kv.pull(3, out = a)\n", |
| "print(a.asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "ee5ff39d", |
| "metadata": {}, |
| "source": [ |
| "`[[ 2. 2. 2.],[ 2. 2. 2.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "## Push, Aggregate, and Update\n", |
| "\n", |
| "For any key that has been initialized, you can push a new value with the same shape to the key:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "f33b4250", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "kv.push(3, mx.nd.ones(shape)*8)\n", |
| "kv.pull(3, out = a) # pull out the value\n", |
| "print(a.asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "54054476", |
| "metadata": {}, |
| "source": [ |
| "`[[ 8. 8. 8.],[ 8. 8. 8.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "The data for pushing can be stored on any device. Furthermore, you can push multiple\n", |
| "values into the same key, where KVStore will first sum all of these\n", |
| "values and then push the aggregated value. Here we will just demonstrate pushing a list of values on CPU.\n", |
| "Please note summation only happens if the value list is longer than one" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "32bf7788", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "contexts = [mx.cpu(i) for i in range(4)]\n", |
| "b = [mx.nd.ones(shape, ctx) for ctx in contexts]\n", |
| "kv.push(3, b)\n", |
| "kv.pull(3, out = a)\n", |
| "print(a.asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "5cf424eb", |
| "metadata": {}, |
| "source": [ |
| "`[[ 4. 4. 4.],[ 4. 4. 4.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "For each push, KVStore combines the pushed value with the value stored using an\n", |
| "`updater`. The default updater is `ASSIGN`. You can replace the default to\n", |
| "control how data is merged:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "bbd1ff4d", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "def update(key, input, stored):\n", |
| " print(\"update on key: %d\" % key)\n", |
| " stored += input * 2\n", |
| "kv._set_updater(update)\n", |
| "kv.pull(3, out=a)\n", |
| "print(a.asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "236fc606", |
| "metadata": {}, |
| "source": [ |
| "`[[ 4. 4. 4.],[ 4. 4. 4.]]`<!--notebook-skip-line-->" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "88ec9860", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "kv.push(3, mx.nd.ones(shape))\n", |
| "kv.pull(3, out=a)\n", |
| "print(a.asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "781251dd", |
| "metadata": {}, |
| "source": [ |
| "`update on key: 3`<!--notebook-skip-line-->\n", |
| "\n", |
| "`[[ 6. 6. 6.],[ 6. 6. 6.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "\n", |
| "## Pull\n", |
| "\n", |
| "You've already seen how to pull a single key-value pair. Similarly, to push, you can\n", |
| "pull the value onto several devices with a single call:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "5a9ec0be", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "b = [mx.nd.ones(shape, ctx) for ctx in contexts]\n", |
| "kv.pull(3, out = b)\n", |
| "print(b[1].asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "54f6b870", |
| "metadata": {}, |
| "source": [ |
| "`[ 6. 6. 6.]],[[ 6. 6. 6.]`<!--notebook-skip-line-->\n", |
| "\n", |
| "## Handle a List of Key-Value Pairs\n", |
| "\n", |
| "All operations introduced so far involve a single key. KVStore also provides\n", |
| "an interface for a list of key-value pairs.\n", |
| "\n", |
| "For a single device:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "8b2cb6e9", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "keys = [5, 7, 9]\n", |
| "kv.init(keys, [mx.nd.ones(shape)]*len(keys))\n", |
| "kv.push(keys, [mx.nd.ones(shape)]*len(keys))\n", |
| "b = [mx.nd.zeros(shape)]*len(keys)\n", |
| "kv.pull(keys, out = b)\n", |
| "print(b[1].asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "4476b6c6", |
| "metadata": {}, |
| "source": [ |
| "`update on key: 5`<!--notebook-skip-line-->\n", |
| "\n", |
| "`update on key: 7`<!--notebook-skip-line-->\n", |
| "\n", |
| "`update on key: 9`<!--notebook-skip-line-->\n", |
| "\n", |
| "`[[ 3. 3. 3.],[ 3. 3. 3.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "For multiple devices:" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "526f512e", |
| "metadata": {}, |
| "source": [ |
| "```python\n", |
| "b = [[mx.nd.ones(shape, ctx) for ctx in contexts]] * len(keys)\n", |
| "kv.push(keys, b)\n", |
| "kv.pull(keys, out = b)\n", |
| "print(b[1][1].asnumpy())\n", |
| "```\n" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "b35c889e", |
| "metadata": {}, |
| "source": [ |
| "`update on key: 5`<!--notebook-skip-line-->\n", |
| "\n", |
| "`update on key: 7`<!--notebook-skip-line-->\n", |
| "\n", |
| "`update on key: 9`<!--notebook-skip-line-->\n", |
| "\n", |
| "`[[ 11. 11. 11.],[ 11. 11. 11.]]`<!--notebook-skip-line-->\n", |
| "\n", |
| "## Run on Multiple Machines\n", |
| "Based on parameter server, the `updater` runs on the server nodes.\n", |
| "When the distributed version is ready, we will update this section.\n", |
| "\n", |
| "\n", |
| "<!-- ## How to Choose Between APIs -->\n", |
| "\n", |
| "<!-- You can mix APIs as much as you like. Here are some guidelines -->\n", |
| "<!-- * Use the Symbolic API and a coarse-grained operator to create an established structure. -->\n", |
| "<!-- * Use a fine-grained operator to extend parts of a more flexible symbolic graph. -->\n", |
| "<!-- * Do some dynamic NDArray tricks, which are even more flexible, between the calls of forward and backward executors. -->\n", |
| "\n", |
| "<!-- Different approaches offer you different levels of flexibility and -->\n", |
| "<!-- efficiency. Normally, you do not need to be flexible in all parts of the -->\n", |
| "<!-- network, so use the parts optimized for speed, and compose it -->\n", |
| "<!-- flexibly with a fine-grained operator or a dynamic NDArray. Such a -->\n", |
| "<!-- mixture allows you to build the deep learning architecture both efficiently and -->\n", |
| "<!-- flexibly as your choice. -->\n", |
| "\n", |
| "## Next Steps\n", |
| "* [MXNet tutorials index](/api/python/docs/tutorials/)\n", |
| "\n", |
| "<!-- INSERT SOURCE DOWNLOAD BUTTONS -->" |
| ] |
| } |
| ], |
| "metadata": { |
| "language_info": { |
| "name": "python" |
| } |
| }, |
| "nbformat": 4, |
| "nbformat_minor": 5 |
| } |