| { |
| "cells": [ |
| { |
| "cell_type": "markdown", |
| "id": "0bf27ff0", |
| "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": "code", |
| "execution_count": 1, |
| "id": "643686b0", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "[[2. 2. 2.]\n", |
| " [2. 2. 2.]]\n" |
| ] |
| }, |
| { |
| "name": "stderr", |
| "output_type": "stream", |
| "text": [ |
| "[04:51:19] /work/mxnet/src/storage/storage.cc:202: Using Pooled (Naive) StorageManager for CPU\n", |
| "[04:51:19] /work/mxnet/src/storage/storage.cc:202: Using Pooled (Naive) StorageManager for CPU_PINNED\n" |
| ] |
| } |
| ], |
| "source": [ |
| "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.np.ones(shape)*2)\n", |
| "a = mx.np.zeros(shape)\n", |
| "kv.pull(3, out = a)\n", |
| "print(a.asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "a081dfe6", |
| "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": "code", |
| "execution_count": 2, |
| "id": "00196413", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "[[8. 8. 8.]\n", |
| " [8. 8. 8.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "kv.push(3, mx.np.ones(shape)*8)\n", |
| "kv.pull(3, out = a) # pull out the value\n", |
| "print(a.asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "bd8e1c88", |
| "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": "code", |
| "execution_count": 3, |
| "id": "8888c0f6", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "[[4. 4. 4.]\n", |
| " [4. 4. 4.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "devices = [mx.cpu(i) for i in range(4)]\n", |
| "b = [mx.np.ones(shape=shape, device=device) for device in devices]\n", |
| "kv.push(3, b)\n", |
| "kv.pull(3, out = a)\n", |
| "print(a.asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "48e07b9b", |
| "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": "code", |
| "execution_count": 4, |
| "id": "cdccd1b0", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "[[4. 4. 4.]\n", |
| " [4. 4. 4.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "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())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "c5e528df", |
| "metadata": {}, |
| "source": [ |
| "`[[ 4. 4. 4.],[ 4. 4. 4.]]`<!--notebook-skip-line-->" |
| ] |
| }, |
| { |
| "cell_type": "code", |
| "execution_count": 5, |
| "id": "da68605b", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "update on key: 3\n", |
| "[[6. 6. 6.]\n", |
| " [6. 6. 6.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "kv.push(3, mx.np.ones(shape))\n", |
| "kv.pull(3, out=a)\n", |
| "print(a.asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "b02146d6", |
| "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": "code", |
| "execution_count": 6, |
| "id": "489bae6a", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "[[6. 6. 6.]\n", |
| " [6. 6. 6.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "b = [mx.np.ones(shape=shape, device=device) for device in devices]\n", |
| "kv.pull(3, out = b)\n", |
| "print(b[1].asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "84ed5225", |
| "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": "code", |
| "execution_count": 7, |
| "id": "818fc728", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "update on key: 5\n", |
| "update on key: 7\n", |
| "update on key: 9\n", |
| "[[3. 3. 3.]\n", |
| " [3. 3. 3.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "keys = [5, 7, 9]\n", |
| "kv.init(keys, [mx.np.ones(shape)]*len(keys))\n", |
| "kv.push(keys, [mx.np.ones(shape)]*len(keys))\n", |
| "b = [mx.np.zeros(shape)]*len(keys)\n", |
| "kv.pull(keys, out = b)\n", |
| "print(b[1].asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "2586695d", |
| "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": "code", |
| "execution_count": 8, |
| "id": "faf0ce40", |
| "metadata": {}, |
| "outputs": [ |
| { |
| "name": "stdout", |
| "output_type": "stream", |
| "text": [ |
| "update on key: 5\n", |
| "update on key: 7\n", |
| "update on key: 9\n", |
| "[[11. 11. 11.]\n", |
| " [11. 11. 11.]]\n" |
| ] |
| } |
| ], |
| "source": [ |
| "b = [[mx.np.ones(shape=shape, device=device) for device in devices]] * len(keys)\n", |
| "kv.push(keys, b)\n", |
| "kv.pull(keys, out = b)\n", |
| "print(b[1][1].asnumpy())" |
| ] |
| }, |
| { |
| "cell_type": "markdown", |
| "id": "d212fa48", |
| "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](../../index.rst)\n", |
| "\n", |
| "<!-- INSERT SOURCE DOWNLOAD BUTTONS -->" |
| ] |
| } |
| ], |
| "metadata": { |
| "language_info": { |
| "name": "python" |
| } |
| }, |
| "nbformat": 4, |
| "nbformat_minor": 5 |
| } |