Apache MXNet - KVStore and Visualization

This chapter explores two key aspects of Apache MXNet: the KVStore for distributed training and data synchronization, and network visualization techniques.

KVStore Package

The Key-Value Store (KVStore) is a fundamental component in MXNet, crucial for training models across multiple devices (GPUs/CPUs) and multiple machines. It facilitates the communication and synchronization of model parameters (weights and biases) during training.

Key concepts of KVStore operation:

Each piece of shared data in the KVStore is identified by a unique ‘key’ (e.g., an integer or string) and holds a ‘value’ (typically an NDArray).
In neural network training, each parameter array (e.g., weights of a layer) is assigned a key. Its value is the actual NDArray containing the parameter’s data.
Worker nodes (devices performing computation) ‘push’ gradients to the KVStore after processing a batch of data. They then ‘pull’ updated parameters from the KVStore before processing the next batch.

Essentially, the KVStore acts as a centralized or distributed repository enabling efficient data sharing and synchronization among various computational units.

Data Push and Pull Operations

Imagine the KVStore as a shared memory object accessible by different devices. Each device can ‘push’ (write) data into it and ‘pull’ (read) data from it.

Let’s walk through the basic implementation steps:

Implementation Steps

Initialization: First, we initialize key-value pairs in the KVStore. Here, we’ll initialize a key (e.g., integer 3) with an NDArray and then pull this value out.

import mxnet as mx
from mxnet import nd

# Create a local KVStore (for single-machine, multi-device scenarios)
kv = mx.kv.create('local')

shape = (3, 3)
key = 3

# Initialize key '3' with a 3x3 NDArray of 2s
kv.init(key, nd.ones(shape) * 2)

# Create an empty NDArray to pull the value into
a = nd.zeros(shape)
kv.pull(key, out=a)

print("Initialized value:")
print(a.asnumpy())

Output:

Initialized value:
[[2. 2. 2.]
 [2. 2. 2.]
 [2. 2. 2.]]

Push, Aggregate, and Update: After initialization, we can push new values to a key. The KVStore can aggregate multiple pushed values or update existing values based on a defined updater function.

Pushing a new value (default behavior is to overwrite):

kv.push(key, nd.ones(shape) * 8)
kv.pull(key, out=a) # Pull the updated value

print("Value after push:")
print(a.asnumpy())

Output:

Value after push:
[[8. 8. 8.]
 [8. 8. 8.]
 [8. 8. 8.]]

Data for pushing can reside on any device (CPU or GPU). If multiple values are pushed to the same key from different sources simultaneously (e.g., gradients from multiple GPUs), the KVStore typically sums these values before applying an update rule. Let’s simulate pushing from multiple contexts:

# Initialize key '3' again for this example, perhaps to zero or a known state.
kv.init(key, nd.zeros(shape)) # Re-initialize to see aggregation clearly

num_devices = 4
# contexts = [mx.gpu(i) for i in range(num_devices)] # If GPUs are available
contexts = [mx.cpu(i) for i in range(num_devices)] # Using CPUs for broad compatibility

# Each device has a piece of data to push (e.g., gradients)
# For this example, let's assume each pushes an array of ones.
# The KVStore will sum these by default if pushed in a list for a single key.
arrays_to_push = [nd.ones(shape, ctx=contexts[i]) for i in range(num_devices)]
kv.push(key, arrays_to_push)
kv.pull(key, out=a)

print("Value after pushing from multiple sources (summed):")
print(a.asnumpy())

Output:

Value after pushing from multiple sources (summed):
[[4. 4. 4.]
 [4. 4. 4.]
 [4. 4. 4.]]

Custom Updater: By default, pushing a single NDArray overwrites the existing value (or uses a predefined aggregation strategy like sum if a list is pushed). We can define a custom updater function to control how new values modify existing ones. The updater is called during a push operation.

# Define a custom updater
def custom_updater(key_id, input_value, stored_value):
    print(f"Updater called on key: {key_id}")
    # Example update rule: stored_value = stored_value + input_value * 2
    stored_value[:] += input_value * 2 # Update in-place

kv.set_updater(custom_updater)

# Current value in 'a' and KVStore for key '3' is [[4., 4., 4.]]
# Let's push a new value
kv.push(key, nd.ones(shape))
# The updater will be called: stored_value (4) += input_value (1) * 2 => 4 + 2 = 6

kv.pull(key, out=a)
print("Value after push with custom updater:")
print(a.asnumpy())

Output (the print from updater might appear before the final print):

Updater called on key: 3
Value after push with custom updater:
[[6. 6. 6.]
 [6. 6. 6.]
 [6. 6. 6.]]

Pull: Similar to push, values can be pulled onto multiple devices in a single call. The out argument can be a list of NDArrays, each on a different context.

b_list = [nd.zeros(shape, ctx=contexts[i]) for i in range(num_devices)]
kv.pull(key, out=b_list)

print("Value pulled onto one of the devices:")
print(b_list[1].asnumpy()) # Displaying the value on the second device/context

Output:

Value pulled onto one of the devices:
[[6. 6. 6.]
 [6. 6. 6.]
 [6. 6. 6.]]

Complete Implementation Example (Consolidated)

import mxnet as mx
from mxnet import nd

kv = mx.kv.create('local')
shape = (3, 3)
key = 3 # Using a single key for simplicity in this example

# 1. Initialization
print("1. Initialization")
kv.init(key, nd.ones(shape) * 2)
a = nd.zeros(shape)
kv.pull(key, out=a)
print(a.asnumpy())

# 2. Push (overwrite)
print("\n2. Push (overwrite)")
kv.push(key, nd.ones(shape) * 8)
kv.pull(key, out=a)
print(a.asnumpy())

# 3. Push from multiple sources (aggregation - sum by default)
print("\n3. Push from multiple sources (aggregation)")
# Re-initialize for clarity of aggregation, or ensure previous value is known.
kv.init(key, nd.zeros(shape)) # Initialize to zeros
num_devices = 4
contexts = [mx.cpu(i) for i in range(num_devices)]
arrays_to_push = [nd.ones(shape, ctx=c) for c in contexts]
kv.push(key, arrays_to_push)
kv.pull(key, out=a)
print(a.asnumpy())

# 4. Custom Updater
print("\n4. Custom Updater")
def custom_updater_example(key_id, input_value, stored_value):
    print(f"Updater logic: key {key_id}, input shape {input_value.shape}, stored shape {stored_value.shape}")
    stored_value[:] += input_value * 2 # Modify stored_value in-place

kv.set_updater(custom_updater_example)
# 'a' currently holds [[4., 4., 4.]] from the previous step.
kv.push(key, nd.ones(shape)) # input_value is ones(shape)
# Updater: stored_value (4) becomes 4 + (1*2) = 6
kv.pull(key, out=a)
print(a.asnumpy())

# 5. Pull to multiple devices
print("\n5. Pull to multiple devices")
b_list = [nd.zeros(shape, ctx=c) for c in contexts]
kv.pull(key, out=b_list)
print(b_list[1].asnumpy()) # Value on one of the devices

Handling Multiple Key-Value Pairs

KVStore efficiently handles operations on lists of key-value pairs.

For a single device

Example for initializing, pushing, and pulling multiple keys on a single device context:

keys = [5, 7, 9]
shape = (2,2) # Using a different shape for variety

# Initialize multiple keys
kv.init(keys, [nd.ones(shape) for _ in keys])

# Push new values to these keys (using the custom_updater from before if still set)
# To avoid unintended effects from previous custom_updater, let's use default updater for this section
# For default ASSIGN behavior with individual pushes:
kv.set_updater(None) # Or, more explicitly, mx.optimizer.Optimizer.get_updater(mx.optimizer.create('sgd')) would yield a no-op essentially for basic push/pull if no training.
                     # For simple overwrite, push new values. If custom updater still active, it will apply.
                     # Let's assume default behavior or re-init KVStore for clean state.
# For this example, let's assume a fresh KV or one where the updater will simply add.
# If using the custom_updater (stored += input * 2):
# init_val = 1. pushed_val = 1. result = 1 + 1*2 = 3.
kv.push(keys, [nd.ones(shape) for _ in keys])

outputs_list = [nd.zeros(shape) for _ in keys]
kv.pull(keys, out=outputs_list)

print("\nValue for one of the keys (e.g., key 7) after multi-key operations:")
# The index in outputs_list corresponds to the index in keys list
print(outputs_list[1].asnumpy()) # Corresponds to key 7

Output (assuming the custom_updater stored_val[:] += input_val * 2 was active, and init was 1, push was 1):

Updater logic: key 5, input shape (2,2), stored shape (2,2)
Updater logic: key 7, input shape (2,2), stored shape (2,2)
Updater logic: key 9, input shape (2,2), stored shape (2,2)

Value for one of the keys (e.g., key 7) after multi-key operations:
[[3. 3.]
 [3. 3.]]

For multiple devices

Pushing data for multiple keys, where each key’s data might be sharded or come from multiple devices:

keys = [10, 11, 12] # New set of keys
shape = (2,2)
num_devices = 2
contexts = [mx.cpu(i) for i in range(num_devices)]

# Initialize keys
kv.init(keys, [nd.ones(shape) * 0.5 for _ in keys]) # Initial value 0.5

# Data to push: for each key, a list of NDArrays, one from each device
# E.g., key 10 gets [data_dev0_key10, data_dev1_key10]
# If custom_updater (val += input*2) is active, and sum aggregation occurs first:
# For key 10: dev0 pushes 1, dev1 pushes 1. Aggregated input = 2.
# Stored = 0.5. New_stored = 0.5 + 2*2 = 4.5
data_to_push_multi_device = []
for _ in keys:
    key_data_from_devices = [nd.ones(shape, ctx=c) for c in contexts]
    data_to_push_multi_device.append(key_data_from_devices)

kv.push(keys, data_to_push_multi_device)

# Pull to multiple devices (or a list of NDArrays on one device)
pulled_data_multi_device = [[nd.zeros(shape, ctx=c) for c in contexts] for _ in keys]
kv.pull(keys, out=pulled_data_multi_device)

print("\nValue for one key (e.g., key 11) on one device (e.g., device 1) after multi-device push:")
print(pulled_data_multi_device[1][1].asnumpy()) # Data for key 11, on context 1

Output (assuming the custom_updater stored_val[:] += input_val * 2 and sum aggregation of inputs per key):

Updater logic: key 10, input shape (2,2), stored shape (2,2)
Updater logic: key 11, input shape (2,2), stored shape (2,2)
Updater logic: key 12, input shape (2,2), stored shape (2,2)

Value for one key (e.g., key 11) on one device (e.g., device 1) after multi-device push:
[[4.5 4.5]
 [4.5 4.5]]

Network Visualization

MXNet’s visualization utilities, such as mxnet.viz.plot_network, help in understanding the architecture of a neural network by rendering it as a computation graph. This graph displays layers or operations as nodes and the flow of data between them as edges. While this typically requires the Graphviz library to be installed and produces an image, the underlying structure can also be understood by examining the symbolic definition of the network.

Visualizing a Neural Network

We can use mx.viz.plot_network to visualize a network defined symbolically. Prerequisites:

A Python environment (e.g., Jupyter Notebook or a script).
Graphviz library (both the core library and the Python graphviz package). You can typically install it via pip install graphviz and system package managers (e.g., apt-get install graphviz on Debian/Ubuntu).

Implementation Example

Let’s define a simple network for linear matrix factorization and describe its structure, as if visualized.

import mxnet as mx
from mxnet import sym

# Define symbolic variables for inputs
user = sym.Variable('user')
item = sym.Variable('item')
score = sym.Variable('score') # Label for training

# Hyperparameters (dummy dimensions for example)
k = 64  # Dimensionality of latent factors
max_user_id = 1000 # Number of unique users
max_item_id = 500  # Number of unique items

# User embedding layer
user_embedding = sym.Embedding(data=user, input_dim=max_user_id, output_dim=k, name='user_embedding')

# Item embedding layer
item_embedding = sym.Embedding(data=item, input_dim=max_item_id, output_dim=k, name='item_embedding')

# Predict score via inner product of user and item embeddings
pred = user_embedding * item_embedding
pred = sym.sum_axis(data=pred, axis=1, name='sum_ reducción') # Corrected name for clarity
pred = sym.Flatten(data=pred, name='flatten_output')

# Define the loss function (e.g., Linear Regression Output for mean squared error)
# This also makes 'score' the label for this network output
network_symbol = sym.LinearRegressionOutput(data=pred, label=score, name='lro')

# To actually plot (if Graphviz is installed):
# viz_graph = mx.viz.plot_network(network_symbol)
# viz_graph.render("matrix_factorization_network", view=False) # Saves to PDF

print("Textual description of the network symbol:")
print("Inputs: 'user', 'item' (features); 'score' (label).")
print(f"1. 'user' -> Embedding (name='user_embedding', input_dim={max_user_id}, output_dim={k})")
print(f"2. 'item' -> Embedding (name='item_embedding', input_dim={max_item_id}, output_dim={k})")
print("3. Outputs of user_embedding and item_embedding -> Element-wise Multiplication")
print("4. Result -> Sum along axis 1 (name='sum_reducción')")
print("5. Result -> Flatten (name='flatten_output')")
print("6. Result (prediction), 'score' (label) -> LinearRegressionOutput (name='lro') for loss calculation.")
print("This network structure represents a common matrix factorization model.")

# For more detailed information on specific layers and their connections,
# one could iterate through network_symbol.get_internals() or analyze the JSON representation:
# print(network_symbol.tojson())

Running mx.viz.plot_network would generate a diagram showing these variables and operations as nodes, with arrows indicating the data flow from inputs through embeddings, multiplication, sum, flatten, and finally to the loss output layer. This visual representation is invaluable for debugging and understanding complex model architectures.