Asynchronous and non-blocking nature of Apache MXNet¶

Instead of using NumPy arrays Apache MXNet offers its own array implementation named NDArray. NDArray API was intentionally designed to be similar to NumPy, but there are differences.

One key difference is in the way calculations are executed. Every NDArray manipulation in Apache MXNet is done in asynchronous, non-blocking way. That means, that when we write code like c = a * b, where both a and b are NDArrays, the function is pushed to the Execution Engine, which starts the calculation. The function immediately returns back, and the user thread can continue execution, despite the fact that the calculation may not have been completed yet.

Execution Engine builds the computation graph which may reorder or combine some calculations, but it honors dependency order: if there are other manipulation with c done later in the code, the Execution Engine will start doing them once the result of c is available. We don’t need to write callbacks to start execution of subsequent code - the Execution Engine is going to do it for us.

To get the result of the computation we only need to access the resulting variable, and the flow of the code will be blocked until the computation results are assigned to the resulting variable. This behavior allows to increase code performance while still supporting imperative programming mode.

Refer to the intro tutorial to NDArray, if you are new to Apache MXNet and would like to learn more how to manipulate NDArrays.

Converting NDArray to NumPy Array blocks calculation¶

Many people are familiar with NumPy and flexible doing tensor manipulations using it. NDArray API offers a convinient .asnumpy() method to cast nd.array to np.array. However, by doing this cast and using np.array for calculation, we cannot use all the goodness of Execution Engine. All manipulations done on np.array are blocking. Moreover, the cast to np.array itself is a blocking operation (same as .asscalar(), .wait_to_read() and .waitall()).

That means that if we have a long computation graph and, at some point, we want to cast the result to np.array, it may feel like the casting takes a lot of time. But what really takes this time is Execution Engine, which finishes all the async calculations we have pushed into it to get the final result, which then will be converted to np.array.

Because of the blocking nature of .asnumpy() method, using it reduces the execution performance, especially if the calculations are done on GPU: Apache MXNet has to copy data from GPU to CPU to return np.array.

The best solution is to make manipulations directly on NDArrays by methods provided in NDArray API.

NumPy operators vs. NDArray operators¶

Despite the fact that NDArray API was specifically designed to be similar to NumPy, sometimes it is not easy to replace existing NumPy computations. The main reason is that not all operators, that are available in NumPy, are available in NDArray API. The list of currently available operators is available on NDArray class page.

If a required operator is missing from NDArray API, there are few things you can do.

from mxnet import nd
import numpy as np

# NumPy has full_like() operator 
np_y = np.full_like(a=np.arange(6, dtype=int), fill_value=10)

# NDArray doesn't have it, but we can replace it with
# creating an array of ones and then multiplying by fill_value
nd_y = nd.ones(shape=(6,)) * 10

# To compare results we had to convert NDArray to NumPy
# But this is okay for that particular case
np.array_equal(np_y, nd_y.asnumpy())

True 

def pad_array(data, max_length):
    # expand dimensions to 4, because nd.pad can work only with 4 dims
    data_expanded = data.reshape(1, 1, 1, data.shape[0])
    
    # pad all 4 dimensions with constant value of 0
    data_padded = nd.pad(data_expanded,
                             mode='constant',
                             pad_width=[0, 0, 0, 0, 0, 0, 0, max_length - data.shape[0]],
                             constant_value=0)
    
    # remove temporary dimensions 
    data_reshaped_back = data_padded.reshape(max_length)
    return data_reshaped_back

pad_array(nd.array([1, 2, 3]), max_length=10)

[ 1.  2.  3.  0.  0.  0.  0.  0.  0.  0.] 

 

How to minimize the impact of blocking calls¶

There are cases, when you have to use either .asnumpy() or .asscalar() methods. As it is explained before, this will force Apache MXNet to block the execution until the result can be retrieved. One common use case is printing a metric or a value of a loss function.

You can minimize the impact of a blocking call by calling .asnumpy() or .asscalar() in the moment, when you think the calculation of this value is already done. In the example below, we introduce the LossBuffer class. It is used to cache the previous value of a loss function. By doing so, we delay printing by one iteration in hope that the Execution Engine would finish the previous iteration and blocking time would be minimized.

from __future__ import print_function

import mxnet as mx
from mxnet import gluon, nd, autograd
from mxnet.ndarray import NDArray
from mxnet.gluon import HybridBlock
import numpy as np

class LossBuffer(object):
    """
    Simple buffer for storing loss value
    """
    def __init__(self):
        self._loss = None

    def new_loss(self, loss):
        ret = self._loss
        self._loss = loss
        return ret

    @property
    def loss(self):
        return self._loss


net = gluon.nn.Dense(10)
ce = gluon.loss.SoftmaxCELoss()
net.initialize()

data = nd.random.uniform(shape=(1024, 100))
label = nd.array(np.random.randint(0, 10, (1024,)), dtype='int32')
train_dataset = gluon.data.ArrayDataset(data, label)
train_data = gluon.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd')
loss_buffer = LossBuffer()

for data, label in train_data:
    with autograd.record():
        out = net(data)
        # This call saves new loss and returns previous loss
        prev_loss = loss_buffer.new_loss(ce(out, label))
        
    loss_buffer.loss.backward()
    trainer.step(data.shape[0])
    
    if prev_loss is not None:
        print("Loss: {}".format(np.mean(prev_loss.asnumpy())))

Loss: 2.310760974884033 

Loss: 2.334498643875122 

Loss: 2.3244147300720215 

Loss: 2.332686424255371 

Loss: 2.321366310119629 

Loss: 2.3236165046691895 

Loss: 2.3178648948669434 

Conclusion¶

For performance reasons, it is better to use native NDArray API methods and avoid using NumPy altogether. In case when you must use NumPy, you can use convenient method .asnumpy() on NDArray to get NumPy representation. By doing so, you block the whole computational process, and force data to be synced between CPU and GPU. If it is a necessary evil to do that, try to minimize the blocking time by calling .asnumpy() in time, when you expect the value to be already computed.