Apache MXNet - NDArray

Apache MXNet - NDArray

In this chapter, we will explore NDArray, Apache MXNet’s primary tool for handling numerical data in a multi-dimensional array format, similar to NumPy arrays but with added capabilities for deep learning and GPU computation.

Handling Data with NDArray

Let’s begin by understanding how to create and manipulate data using NDArray. You’ll need the following prerequisites:

Prerequisites

To effectively work with NDArrays and follow the examples in this chapter, ensure you have:

• Apache MXNet (version 1.6 or later recommended) installed in your Python environment. You can typically install it using pip: pip install mxnet (or a variant with GPU support like pip install mxnet-cu118 depending on your CUDA version).
• Python 3.7 or later.

Basic NDArray Creation

Let’s dive into some examples to understand basic NDArray functionality.

First, import the MXNet library and its NDArray module. A common convention is to import ndarray as nd:

import mxnet as mx
from mxnet import nd

Now, let’s explore various ways to create NDArrays:

Creating a 1-D Array from a Python List

Example

x = nd.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
print(x)

Output

The output shows a 1-D array:

[ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10.]
<ndarray shape=(10,) cpu(0)>

Creating a 2-D Array from a Python List

Example

y = nd.array([
    [1, 2, 3, 4, 5],
    [6, 7, 8, 9, 10],
    [11, 12, 13, 14, 15]
])
print(y)

Output

The output shows a 2-D array (a matrix):

[[ 1.  2.  3.  4.  5.]
 [ 6.  7.  8.  9. 10.]
 [11. 12. 13. 14. 15.]]
<ndarray shape=(3, 5) cpu(0)>

Creating NDArrays with Specific Initializers

MXNet provides functions to create NDArrays with specific shapes and initial content:

Example: Using nd.empty, nd.full

nd.empty(shape, ctx) creates an array with the specified shape on the given context, but its entries are uninitialized (containing arbitrary data from memory). It’s slightly faster if you plan to fill it immediately, but use with caution.

nd.full(shape, value, ctx) creates an array of a given shape and context, filled entirely with a specified scalar value.

a = nd.empty((2, 3), ctx=mx.cpu())  # Uninitialized 2x3 array on CPU
print("Uninitialized array (a):")
print(a)

b = nd.full((2, 3), 8.0, ctx=mx.cpu()) # 2x3 array filled with 8.0 on CPU
print("\nArray filled with 8.0 (b):")
print(b)

Output

The output will be similar to this (exact values for nd.empty may vary each time):

Uninitialized array (a):
[[0.000e+00 0.000e+00 0.000e+00]
 [0.000e+00 2.887e-42 0.000e+00]]
<ndarray shape=(2, 3) cpu(0)>

Array filled with 8.0 (b):
[[8. 8. 8.]
 [8. 8. 8.]]
<ndarray shape=(2, 3) cpu(0)>

Creating an Array of Zeros

Use nd.zeros(shape, ctx) to create an array filled with zeros on the specified context.

Example

z = nd.zeros((3, 4), ctx=mx.cpu())
print(z)

Output

The output is a 3x4 matrix of zeros:

[[0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]]
<ndarray shape=(3, 4) cpu(0)>

Creating an Array of Ones

Use nd.ones(shape, ctx) to create an array filled with ones on the specified context.

Example

o = nd.ones((2, 5), ctx=mx.cpu())
print(o)

Output

The output is a 2x5 matrix of ones:

[[1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1.]]
<ndarray shape=(2, 5) cpu(0)>

Creating an Array with Random Values

MXNet’s nd.random submodule offers various distributions. For example, nd.random.normal(loc, scale, shape, ctx) samples from a normal (Gaussian) distribution on a specified context.

Example

# Sample from a normal distribution with mean 0 and standard deviation 1
r = nd.random.normal(loc=0, scale=1, shape=(3, 4), ctx=mx.cpu())
print(r)

Output

The output will be a 3x4 array of random numbers (values will vary on each run):

[[ 1.2673576  -2.0345826  -0.32537818 -1.4583491 ]
 [-0.11176403  1.3606371  -0.7889914  -0.17639421]
 [-0.2532185  -0.42614475 -0.12548696  1.4022992 ]]
<ndarray shape=(3, 4) cpu(0)>

Inspecting NDArray Properties

You can inspect key properties of an NDArray:

Example: Get shape, size, data type, and context of the random array r created above.

print("Shape of r:", r.shape)
print("Size of r (total number of elements):", r.size)
print("Data type of r:", r.dtype)
print("Context of r:", r.context)

Output

The output will be:

Shape of r: (3, 4)
Size of r (total number of elements): 12
Data type of r: <class 'numpy.float32'>
Context of r: cpu(0)

Note: MXNet defaults to 32-bit floating-point numbers (float32). The dtype is often reported as a NumPy dtype due to MXNet’s interoperability. The .context attribute shows where the array’s data resides (e.g., cpu(0) for the first CPU, gpu(0) for the first GPU).

NDArray Operations

NDArray supports a comprehensive set of mathematical operations, many of which are similar to NumPy. These operations are typically element-wise unless specified otherwise and are context-aware (i.e., operations usually happen on the context of the input arrays).

Element-wise Addition

Example

import mxnet as mx
from mxnet import nd

# Ensure arrays are on the same context for operations
ctx = mx.cpu()
x = nd.ones((2, 3), ctx=ctx)
y = nd.arange(6, ctx=ctx).reshape((2, 3)) # Creates [[0,1,2],[3,4,5]]

print('x =\n', x)
print('y =\n', y)

result = x + y
print('result (x + y) =\n', result)

Output

The output demonstrates element-wise addition:

x =
[[1. 1. 1.]
 [1. 1. 1.]]
<ndarray shape=(2, 3) cpu(0)>
y =
[[0. 1. 2.]
 [3. 4. 5.]]
<ndarray shape=(2, 3) cpu(0)>
result (x + y) =
[[1. 2. 3.]
 [4. 5. 6.]]
<ndarray shape=(2, 3) cpu(0)>

Element-wise Multiplication

Example

ctx = mx.cpu()
a = nd.array([1, 2, 3, 4], ctx=ctx)
b = nd.array([2, 2, 2, 1], ctx=ctx)

product = a * b
print('product (a * b) =', product)

Output

The output is:

product (a * b) = [2. 4. 6. 4.]
<ndarray shape=(4,) cpu(0)>

Exponentiation

nd.exp(array) computes the element-wise exponential (e^x).

Example

c = nd.array([1, 2, 3], ctx=mx.cpu())
exp_c = nd.exp(c)
print('exp_c =', exp_c)

Output

The output shows e raised to the power of each element:

exp_c = [ 2.7182817  7.389056  20.085537 ]
<ndarray shape=(3,) cpu(0)>

Matrix Multiplication

Use nd.dot(A, B) for matrix multiplication. For 2-D arrays, this computes the standard matrix product. Python 3.5+ also supports the @ operator for matrix multiplication, which is often preferred for its conciseness.

Example

ctx = mx.cpu()
mat1 = nd.arange(6, ctx=ctx).reshape((2, 3)) # 2x3 matrix: [[0,1,2],[3,4,5]]
mat2 = nd.arange(3, ctx=ctx).reshape((3, 1)) # 3x1 matrix: [[0],[1],[2]]

# Using nd.dot
dot_product = nd.dot(mat1, mat2)
print('Using nd.dot:\n', dot_product)

# Using @ operator (requires Python 3.5+)
at_product = mat1 @ mat2
print('\nUsing @ operator:\n', at_product)

Output

Both methods yield the same 2x1 matrix:

Using nd.dot:
[[ 5.]
 [14.]]
<ndarray shape=(2, 1) cpu(0)>

Using @ operator:
[[ 5.]
 [14.]]
<ndarray shape=(2, 1) cpu(0)>

In-place Operations

Standard operations like y = x + y allocate new memory for the result y. In deep learning, this can be inefficient for large arrays or frequent updates (e.g., model parameters during training). MXNet supports in-place operations to modify existing array memory, reducing overhead.

Let’s illustrate memory allocation. Python’s id() function returns an object’s unique identifier (often its memory address in CPython).

ctx = mx.cpu()
x = nd.arange(4, ctx=ctx)
y = nd.ones(4, ctx=ctx)

print('Original y =', y)
print('id(y) before addition:', id(y))

y_new = y + x  # Standard addition allocates new memory for y_new
print('\ny_new after y_new = y + x:', y_new)
print('id(y_new) after addition (new memory):', id(y_new))
print('id(y) (original y is unchanged):', id(y))

Output

Note how id(y_new) is different from id(y), indicating y_new points to a new memory location:

Original y = [1. 1. 1. 1.]
<ndarray shape=(4,) cpu(0)>
id(y) before addition: 140123456789012

y_new after y_new = y + x: [1. 2. 3. 4.]
<ndarray shape=(4,) cpu(0)>
id(y_new) after addition (new memory): 140123456789134
id(y) (original y is unchanged): 140123456789012

(Note: Exact id values will vary.)

To perform operations in-place, you can use slice assignment or functions with an out argument:

1. Using slice assignment [:]:

ctx = mx.cpu()
x = nd.arange(4, ctx=ctx)
y_orig = nd.ones(4, ctx=ctx)
z = nd.zeros_like(x) # Pre-allocate z on the same context as x

print('id(z) before:', id(z))
# Result of y_orig + x is first computed (possibly in new memory),
# then copied into z's existing memory via slice assignment.
z[:] = y_orig + x

print('\nz after z[:] = y_orig + x:', z)
print('id(z) after (should be same):', id(z))

Output

The id(z) remains the same, but a temporary buffer might have been used for y_orig + x before the copy.

id(z) before: 140123456789256

z after z[:] = y_orig + x: [1. 2. 3. 4.]
<ndarray shape=(4,) cpu(0)>
id(z) after (should be same): 140123456789256

2. Using the out argument (most memory-efficient): Many MXNet operators (like nd.add, nd.multiply) accept an out argument to specify the destination array. This avoids temporary allocations for the operation itself.

ctx = mx.cpu()
x = nd.arange(4, ctx=ctx)
y = nd.ones(4, ctx=ctx)
z_out = nd.zeros_like(x) # Pre-allocate z_out

print('id(z_out) before:', id(z_out))
nd.add(x, y, out=z_out) # Efficient in-place capable addition

print('\nz_out after nd.add(x, y, out=z_out):', z_out)
print('id(z_out) after (should be same):', id(z_out))

Output

The id(z_out) remains unchanged, and ideally, no intermediate buffer is needed for the addition operation itself.

id(z_out) before: 140123456789378

z_out after nd.add(x, y, out=z_out): [1. 2. 3. 4.]
<ndarray shape=(4,) cpu(0)>
id(z_out) after (should be same): 140123456789378

Shorthand in-place operators like x += y or x *= y also modify x directly when supported. For example, z_out += x (after z_out holds x+y) would add x to z_out in-place if z_out’s memory can be reused.

Managing Computation Context (CPU/GPU)

MXNet allows you to specify the computation context (CPU or GPU) for NDArrays. This is crucial for performance, as computations on GPUs are typically much faster for deep learning tasks. By default, NDArrays are created on the CPU unless a different default context is set (mx.default_context()).

You can check for GPU availability: num_gpus = mx.context.num_gpus(). If num_gpus is 0, GPU examples will raise an error or require modification.

Example: Creating an array on a specific CPU, and then copying to a GPU (if available).

import mxnet as mx
from mxnet import nd

# Create an array on the default CPU (often cpu(0))
default_cpu_ctx = mx.cpu() # Or mx.cpu(0)
a_cpu = nd.ones(shape=(3, 3), ctx=default_cpu_ctx)
print("Array on CPU:")
print(a_cpu)

# Check if GPUs are available
num_gpus = mx.context.num_gpus()
print(f"\nNumber of GPUs available: {num_gpus}")

if num_gpus > 0:
    gpu_ctx = mx.gpu(0) # Use the first GPU
    # Copy array to the GPU
    a_gpu = a_cpu.copyto(gpu_ctx)
    print("\nArray copied to GPU:")
    print(a_gpu)

    # Alternative way to get an array on a context:
    # as_in_context returns the array itself if already on target context,
    # otherwise, it copies and returns a new array on the target context.
    b_gpu = nd.ones(shape=(2,2), ctx=gpu_ctx)
    c_gpu = a_cpu.as_in_context(gpu_ctx) # c_gpu is a new array on gpu(0)
    d_gpu = b_gpu.as_in_context(gpu_ctx) # d_gpu is b_gpu itself (no copy)
    print("\nArray c_gpu (from as_in_context on a_cpu):")
    print(c_gpu)
    print(f"id(b_gpu) == id(d_gpu): {id(b_gpu) == id(d_gpu)}")
else:
    print("\nNo GPUs available. Skipping GPU-specific copy examples.")

Output (if a GPU is present):

Array on CPU:
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]
<ndarray shape=(3, 3) cpu(0)>

Number of GPUs available: 1 (output may vary)

Array copied to GPU:
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]
<ndarray shape=(3, 3) gpu(0)>

Array c_gpu (from as_in_context on a_cpu):
[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]
<ndarray shape=(3, 3) gpu(0)>
id(b_gpu) == id(d_gpu): True

Operations involving NDArrays on different contexts (e.g., adding a CPU array to a GPU array) are generally not allowed directly. Data must be explicitly moved to the same context. Transferring data between CPU and GPU incurs overhead, so it’s best practice to minimize these transfers in performance-critical code.

NDArray vs. NumPy Array: Key Differences

While NDArray is designed to be similar to NumPy’s array, there are key distinctions, primarily related to deep learning optimizations:

• GPU Acceleration: NDArray has native support for GPU computation, essential for training deep neural networks efficiently.
• Asynchronous Execution: MXNet operations on NDArrays are often executed asynchronously. When you issue a command like c = a + b, the operation is pushed to MXNet’s execution engine. The Python interpreter might regain control immediately, even before the computation completes. The engine manages dependencies and optimizes the execution order. This allows for parallelism and efficiency, especially with GPUs. Actual results are typically only computed and made available when explicitly requested (e.g., printing, converting to NumPy via .asnumpy(), or saving to disk).
• Automatic Differentiation: NDArray integrates seamlessly with MXNet’s autograd module for automatic calculation of gradients (derivatives), which is fundamental to training neural networks via backpropagation.

To ensure a computation has finished and to get its actual result, you can call .wait_to_read() on an NDArray. Converting it to a NumPy array using .asnumpy() is also a blocking call that waits for completion.

Converting NDArray to and from NumPy Array

MXNet provides seamless conversion between NDArrays and NumPy arrays, facilitating integration with the broader Python scientific computing ecosystem.

NDArray to NumPy: Use the .asnumpy() method.

import numpy as np
from mxnet import nd

nd_arr = nd.array([1, 2, 3], ctx=mx.cpu())
np_arr = nd_arr.asnumpy() # Convert NDArray to NumPy array

print("NDArray:", nd_arr)
print("NumPy array:", np_arr)
print("Type of np_arr:", type(np_arr))

Output:

NDArray: [1. 2. 3.]
<ndarray shape=(3,) cpu(0)>
NumPy array: [1. 2. 3.]
Type of np_arr: <class 'numpy.ndarray'>

NumPy to NDArray: Use nd.array(numpy_array, ctx).

np_data = np.array([[4, 5], [6, 7]])
nd_data = nd.array(np_data, ctx=mx.cpu()) # Convert NumPy array to NDArray on CPU

print("Original NumPy array:\n", np_data)
print("Converted NDArray:\n", nd_data)

Output:

Original NumPy array:
[[4 5]
 [6 7]]
Converted NDArray:
[[4. 5.]
 [6. 7.]]
<ndarray shape=(2, 2) cpu(0)>

Important: Converting an NDArray residing on a GPU to a NumPy array using .asnumpy() will implicitly copy the data from GPU to CPU memory. This operation can be slow and should be used judiciously in performance-critical code sections.

You can often achieve functionality similar to some NumPy functions by combining existing NDArray operators. For example, np.full_like(a, val) (create an array with the same shape as a, filled with val) can be replicated in MXNet:

a_np = np.arange(5, dtype=int)
np_x = np.full_like(a=a_np, fill_value=10)

# MXNet equivalent
ctx = mx.cpu()
a_nd = nd.arange(5, ctx=ctx)
# Create an array of ones with the same shape and context as a_nd, then multiply by 10
nd_x = nd.ones_like(a_nd) * 10
# Alternatively, more directly:
# nd_x_alt = nd.full(a_nd.shape, 10, ctx=a_nd.context, dtype=a_nd.dtype)

print("NumPy full_like result:", np_x)
print("MXNet equivalent result:", nd_x)
print("Are they equal (after converting nd_x to numpy)?", np.array_equal(np_x, nd_x.asnumpy()))

Output:

NumPy full_like result: [10 10 10 10 10]
MXNet equivalent result: [10. 10. 10. 10. 10.]
<ndarray shape=(5,) cpu(0)>
Are they equal (after converting nd_x to numpy)? True

While many MXNet NDArray functions mirror NumPy, names or signatures might differ slightly (e.g., nd.split() vs np.split()). Always consult the official Apache MXNet API documentation for precise usage. For a more direct NumPy-like experience within MXNet, consider using the mxnet.numpy module (usually imported as mx.np), which offers an API that closely follows NumPy’s conventions and often provides better NumPy compatibility.

Minimizing Impact of Blocking Calls in Training Loops

Calls like .asnumpy() or .asscalar() (to get a Python scalar from a single-element NDArray) are blocking: they force MXNet to finish all pending computations for that array and then copy the data from its current context (e.g., GPU) to CPU memory if necessary. In training loops, frequent blocking calls can create performance bottlenecks by serializing execution.

One strategy to mitigate this is to delay these blocking calls until the result is truly needed or, ideally, to overlap them with other independent computations. For instance, when logging training loss, instead of converting the loss to a NumPy array immediately in each iteration, one can buffer the NDArray loss and convert/print the loss from the previous iteration. This allows the current iteration’s computations (including backpropagation on GPU) to potentially overlap with the (blocking, CPU-bound) asnumpy() call from the prior step.

Implementation Example

The following example demonstrates this concept using a simple MXNet Gluon API training loop. The LossBuffer class helps delay the asnumpy() call on the loss value.

import mxnet as mx
from mxnet import gluon, nd, autograd
import numpy as np
import time # For demonstration

class LossBuffer(object):
    """Simple buffer for storing and retrieving loss, delaying asnumpy()."""
    def __init__(self):
        self._loss_nd_current_iter = None # Stores NDArray loss of current iteration
        self._loss_np_previous_iter = None # Stores NumPy loss from previous iteration

    def log_and_buffer(self, new_loss_nd):
        """Converts previous loss to NumPy, buffers new loss as NDArray, returns previous NumPy loss."""
        # Convert previous iteration's loss (if it exists) to NumPy
        if self._loss_nd_current_iter is not None:
            # This asnumpy() call is on the *previous* iteration's loss.
            # It might overlap with the current iteration's GPU work.
            self._loss_np_previous_iter = np.mean(self._loss_nd_current_iter.asnumpy())
        else:
            self._loss_np_previous_iter = None # No previous loss on first call

        self._loss_nd_current_iter = new_loss_nd # Store current loss (still an NDArray)
        return self._loss_np_previous_iter

# Setup: A simple model, loss, and data
ctx = mx.gpu() if mx.context.num_gpus() > 0 else mx.cpu() # Use GPU if available, else CPU
print(f"Training on: {ctx}")

net = gluon.nn.Dense(10, prefix='dense_') # Output layer for 10 classes
net.initialize(ctx=ctx) # Initialize parameters on the chosen context

loss_fn = gluon.loss.SoftmaxCrossEntropyLoss() # Standard loss for classification

# Dummy data parameters
num_samples = 1024
num_features = 100
num_batches = 8
batch_size = num_samples // num_batches
num_epochs = 1 # For a quick demo

# Create dummy data directly on the target context
data = nd.random.uniform(shape=(num_samples, num_features), ctx=ctx)
label = nd.array(np.random.randint(0, 10, (num_samples,)), ctx=ctx, dtype='int32')

train_dataset = gluon.data.ArrayDataset(data, label)
# Using num_workers > 0 can speed up data loading by using separate processes.
# However, for simplicity or if issues arise, num_workers=0 (loads in main process) is safer.
train_data_loader = gluon.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd', optimizer_params={'learning_rate': 0.01})

loss_buffer = LossBuffer()

for epoch in range(num_epochs):
    epoch_start_time = time.time()
    cumulative_loss_sum_for_epoch = 0
    num_samples_in_epoch = 0

    for i, (batch_data, batch_label) in enumerate(train_data_loader):
        # Data might be on CPU from DataLoader if num_workers > 0; ensure it's on target context
        batch_data = batch_data.as_in_context(ctx)
        batch_label = batch_label.as_in_context(ctx)

        with autograd.record(): # Start recording computations for gradient calculation
            output = net(batch_data)
            current_batch_loss_nd = loss_fn(output, batch_label) # Loss is an NDArray on ctx

        # This call saves the new NDArray loss and returns the *previous* iteration's loss as NumPy
        previous_avg_loss_np = loss_buffer.log_and_buffer(current_batch_loss_nd)

        current_batch_loss_nd.backward() # Backpropagate based on current_batch_loss_nd
        trainer.step(batch_data.shape[0]) # Update model parameters

        if previous_avg_loss_np is not None:
            # Print the average loss from the *previous* iteration (already converted)
            print(f"Epoch {epoch}, Batch {i-1}: Prev. Iter. Avg Batch Loss: {previous_avg_loss_np:.4f}")

        # Accumulate current batch's loss for epoch average (on NDArray to avoid blocking)
        # Note: current_batch_loss_nd is typically a vector of losses, one per sample.
        # We sum it here; it will be averaged at the end of epoch.
        cumulative_loss_sum_for_epoch += current_batch_loss_nd.sum()
        num_samples_in_epoch += batch_data.shape[0]

    # After the loop, process the final buffered loss (from the last batch of the epoch)
    final_buffered_avg_loss_np = loss_buffer.log_and_buffer(None) # Clear buffer, get last batch's loss
    if final_buffered_avg_loss_np is not None:
        print(f"Epoch {epoch}, Batch {i}: Prev. Iter. Avg Batch Loss: {final_buffered_avg_loss_np:.4f}")

    # Calculate and print average epoch loss (now we need to block for sum)
    avg_epoch_loss = cumulative_loss_sum_for_epoch.asscalar() / num_samples_in_epoch
    print(f"Epoch {epoch} completed. Avg Epoch Loss: {avg_epoch_loss:.4f}. Time: {time.time() - epoch_start_time:.2f}s.")

Output (will vary due to random data/initialization and context; loss values are illustrative):

Training on: cpu(0) (or gpu(0) if available)
Epoch 0, Batch 0: Prev. Iter. Avg Batch Loss: 2.3512
Epoch 0, Batch 1: Prev. Iter. Avg Batch Loss: 2.3488
...
Epoch 0, Batch 6: Prev. Iter. Avg Batch Loss: 2.3078
Epoch 0, Batch 7: Prev. Iter. Avg Batch Loss: 2.3002
Epoch 0 completed. Avg Epoch Loss: 2.3250. Time: 0.15s.

This technique of staggering the .asnumpy() call helps to overlap this potentially blocking CPU-bound task with the independent GPU computations of the subsequent iteration, thereby improving overall training throughput. For more advanced profiling and performance optimization in MXNet, refer to the official Apache MXNet documentation and tools like the MXNet profiler. Additional resources on NDArray operations and best practices can also be found there.