Setup: Build a model¶

Let’s build a small convolutional neural network that we can use for profiling.

from mxnet import gluon
net = gluon.nn.HybridSequential()
with net.name_scope():
    net.add(gluon.nn.Conv2D(channels=20, kernel_size=5, activation='relu'))
    net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
    net.add(gluon.nn.Conv2D(channels=50, kernel_size=5, activation='relu'))
    net.add(gluon.nn.MaxPool2D(pool_size=2, strides=2))
    net.add(gluon.nn.Flatten())
    net.add(gluon.nn.Dense(512, activation="relu"))
    net.add(gluon.nn.Dense(10))

We need data that we can run through the network for profiling. We’ll use the MNIST dataset.

from mxnet.gluon.data.vision import transforms
train_data = gluon.data.DataLoader(gluon.data.vision.MNIST(train=True).transform_first(transforms.ToTensor()),
                                   batch_size=64, shuffle=True)

Let’s define a method that will run one training iteration given data and label.

# Use GPU if available
try:
    mx.test_utils.list_gpus(); ctx = mx.gpu()
except:
    ctx = mx.cpu()

# Initialize the parameters with random weights
net.collect_params().initialize(mx.init.Xavier(), ctx=ctx)

# Use SGD optimizer
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': .1})

# Softmax Cross Entropy is a frequently used loss function for multi-classs classification
softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()

# A helper function to run one training iteration
def run_training_iteration(data, label):
    
    # Load data and label is the right context
    data = data.as_in_context(ctx)
    label = label.as_in_context(ctx)
    
    # Run the forward pass
    with autograd.record():
        output = net(data)
        loss = softmax_cross_entropy(output, label)
    
    # Run the backward pass
    loss.backward()
    
    # Apply changes to parameters
    trainer.step(data.shape[0])

Starting and stopping the profiler from Python¶

When the first forward pass is run on a network, MXNet does a number of housekeeping tasks including inferring the shapes of various parameters, allocating memory for intermediate and final outputs, etc. For these reasons, profiling the first iteration doesn’t provide accurate results. We will, therefore skip the first iteration.

# Run the first iteration without profiling
itr = iter(train_data)
run_training_iteration(*next(itr))

We’ll run the next iteration with the profiler turned on.

data, label = next(itr)

# Ask the profiler to start recording
profiler.set_state('run')

run_training_iteration(*next(itr))

# Ask the profiler to stop recording after operations have completed
mx.nd.waitall()
profiler.set_state('stop')

Between running and stopping the profiler, you can also pause and resume the profiler using profiler.pause() and profiler.resume() respectively to profile only parts of the code you want to profile.

Starting profiler automatically using environment variable¶

The method described above requires code changes to start and stop the profiler. You can also start the profiler automatically and profile the entire code without any code changes using the MXNET_PROFILER_AUTOSTART environment variable.

MXNet will start the profiler automatically if you run your code with the environment variable MXNET_PROFILER_AUTOSTART set to 1. The profiler output is stored into profile.json in the current directory.

Note that the profiler output could be large depending on your code. It might be helpful to profile only sections of your code using the set_state API described in the previous section.

Increasing granularity of the profiler output¶

MXNet executes computation graphs in ‘bulk mode’ which reduces kernel launch gaps in between symbolic operators for faster execution. This could reduce the granularity of the profiler output. If you need profiling result of every operator, please set the environment variables MXNET_EXEC_BULK_EXEC_INFERENCE and MXNET_EXEC_BULK_EXEC_TRAIN to 0 to disable the bulk execution mode.

Viewing profiler output¶

There are two ways to view the information collected by the profiler. You can either view it in the console or you can view a more graphical version in a browser.

print(profiler.dumps())

profiler.dump()

Further reading¶