Apache Mxnet 简明教程

Apache MXNet - KVStore and Visualization

本章涉及 Python 软件包 KVStore 和可视化。

KVStore package

KV 存储表示键值存储。这是用于多设备训练的关键组件。它很重要,因为通过一个或多个具有参数的 KVStore 服务器在单个计算机或多个计算机上的设备之间传输参数通信。

让我们借助以下几点了解 KVStore 的工作原理:

  1. KVStore 中的每个值都由 keyvalue 表示。

  2. 网络中的每个参数数组都被分配一个 key ,并且该参数数组的权重由 value 引用。

  3. 然后,工作节点在处理一个批次后 push 梯度。在处理一个新批次之前,它们还 pull 更新权重。

简而言之,我们可以说 KVStore 是一个数据共享的地方,每个设备都可以将数据推入和拉出。

Data Push-In and Pull-Out

KVStore 可以被认为是在不同设备(如 GPU 和计算机)之间共享的单个对象,其中每个设备都能够将数据推入和拉出。

以下是由设备遵循以将数据推入和拉出的实施步骤:

Implementation steps

Initialisation − 第一步是初始化值。这里,对于我们的示例,我们将在 KVStrore 中初始化一个元组(int、NDArray),然后将值提取出来 −

import mxnet as mx
kv = mx.kv.create('local') # create a local KVStore.
shape = (3,3)
kv.init(3, mx.nd.ones(shape)*2)
a = mx.nd.zeros(shape)
kv.pull(3, out = a)
print(a.asnumpy())

Output

生成以下输出:

[[2. 2. 2.]
[2. 2. 2.]
[2. 2. 2.]]

Push, Aggregate, and Update − 初始化后,我们可以使用相同的形状将新值推入 KVStore 至键 −

kv.push(3, mx.nd.ones(shape)*8)
kv.pull(3, out = a)
print(a.asnumpy())

Output

输出如下 −

[[8. 8. 8.]
 [8. 8. 8.]
 [8. 8. 8.]]

用于推送的数据可以存储在任何设备(如 GPU 或计算机)上。我们还可以将多个值推入同一键。在这种情况下,KVStore 将首先对所有这些值求和,然后推入聚合值,如下所示: −

contexts = [mx.cpu(i) for i in range(4)]
b = [mx.nd.ones(shape, ctx) for ctx in contexts]
kv.push(3, b)
kv.pull(3, out = a)
print(a.asnumpy())

Output

您将看到以下输出 −

[[4. 4. 4.]
 [4. 4. 4.]
 [4. 4. 4.]]

对于您应用的每个推送,KVStore 都会将已推送的值与已存储的值合并。它将借助于更新器完成。这里,默认更新器为 ASSIGN。

def update(key, input, stored):
   print("update on key: %d" % key)

   stored += input * 2
kv.set_updater(update)
kv.pull(3, out=a)
print(a.asnumpy())

Output

执行以上代码时,应该看到以下输出 −

[[4. 4. 4.]
 [4. 4. 4.]
 [4. 4. 4.]]

Example

kv.push(3, mx.nd.ones(shape))
kv.pull(3, out=a)
print(a.asnumpy())

Output

以下是代码的输出 −

update on key: 3
[[6. 6. 6.]
 [6. 6. 6.]
 [6. 6. 6.]]

Pull − 与推送一样,我们还可以通过一次调用在多个设备上拉取值,如下所示 −

b = [mx.nd.ones(shape, ctx) for ctx in contexts]
kv.pull(3, out = b)
print(b[1].asnumpy())

Output

输出如下 −

[[6. 6. 6.]
 [6. 6. 6.]
 [6. 6. 6.]]

Complete Implementation Example

下面给出完整的实施示例 −

import mxnet as mx
kv = mx.kv.create('local')
shape = (3,3)
kv.init(3, mx.nd.ones(shape)*2)
a = mx.nd.zeros(shape)
kv.pull(3, out = a)
print(a.asnumpy())
kv.push(3, mx.nd.ones(shape)*8)
kv.pull(3, out = a) # pull out the value
print(a.asnumpy())
contexts = [mx.cpu(i) for i in range(4)]
b = [mx.nd.ones(shape, ctx) for ctx in contexts]
kv.push(3, b)
kv.pull(3, out = a)
print(a.asnumpy())
def update(key, input, stored):
   print("update on key: %d" % key)
   stored += input * 2
kv._set_updater(update)
kv.pull(3, out=a)
print(a.asnumpy())
kv.push(3, mx.nd.ones(shape))
kv.pull(3, out=a)
print(a.asnumpy())
b = [mx.nd.ones(shape, ctx) for ctx in contexts]
kv.pull(3, out = b)
print(b[1].asnumpy())

Handling Key-Value Pairs

我们在上面实现的所有操作都涉及单个键,但 KVStore 还提供了一个 a list of key-value pairs 的接口 −

For a single device

以下是一个示例,展示了 KVStore 接口,用于单个设备的一系列键值对 −

keys = [5, 7, 9]
kv.init(keys, [mx.nd.ones(shape)]*len(keys))
kv.push(keys, [mx.nd.ones(shape)]*len(keys))
b = [mx.nd.zeros(shape)]*len(keys)
kv.pull(keys, out = b)
print(b[1].asnumpy())

Output

您将收到以下输出 −

update on key: 5
update on key: 7
update on key: 9
[[3. 3. 3.]
 [3. 3. 3.]
 [3. 3. 3.]]

For multiple device

以下是一个示例,展示了 KVStore 接口,用于多个设备的一系列键值对 −

b = [[mx.nd.ones(shape, ctx) for ctx in contexts]] * len(keys)
kv.push(keys, b)
kv.pull(keys, out = b)
print(b[1][1].asnumpy())

Output

您将看到以下输出 −

update on key: 5
update on key: 7
update on key: 9
[[11. 11. 11.]
 [11. 11. 11.]
 [11. 11. 11.]]

Visualization package

可视化包是 Apache MXNet 包,用于将神经网络 (NN) 表示为包含节点和边的计算图。

Visualize neural network

在下面的示例中,我们将使用 mx.viz.plot_network 可视化神经网络。以下是其先决条件 −

Prerequisites

  1. Jupyter notebook

  2. Graphviz library

Implementation Example

在下面的示例中,我们将可视化用于线性矩阵分解的样本 NN −

import mxnet as mx
user = mx.symbol.Variable('user')
item = mx.symbol.Variable('item')
score = mx.symbol.Variable('score')

# Set the dummy dimensions
k = 64
max_user = 100
max_item = 50

# The user feature lookup
user = mx.symbol.Embedding(data = user, input_dim = max_user, output_dim = k)

# The item feature lookup
item = mx.symbol.Embedding(data = item, input_dim = max_item, output_dim = k)

# predict by the inner product and then do sum
N_net = user * item
N_net = mx.symbol.sum_axis(data = N_net, axis = 1)
N_net = mx.symbol.Flatten(data = N_net)

# Defining the loss layer
N_net = mx.symbol.LinearRegressionOutput(data = N_net, label = score)

# Visualize the network
mx.viz.plot_network(N_net)