Setup: Build a model¶

Let’s build a small convolutional neural network that we can use to demonstrate 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

dataset = gluon.data.vision.MNIST(train=True)
dataset = dataset.transform_first(transforms.ToTensor())
dataloader = gluon.data.DataLoader(dataset, batch_size=64, shuffle=True)

Let’s define a function that will run a single training iteration given data and label.

# Use GPU if available
if mx.context.num_gpus():
    ctx=mx.gpu()
else:
    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': 0.1})

# Softmax Cross Entropy is a frequently used loss function for multi-class 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 representative results for the rest of training. We will, therefore, skip the first iteration.

# Run the first iteration without profiling
itr = iter(dataloader)
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))

# Make sure all operations have completed
mx.nd.waitall()
# Ask the profiler to stop recording
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 the profiler automatically using an 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_PROFILER_AUTOSTART=1 python my_script.py

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 in profile.json inside 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.

When working with networks created using the Gluon API, you will get a more granular profiling outputs if you profile networks that haven’t been hybridized. Operations can appear fused together in the profiling outputs after hybridization, which can make debugging tricky.

Viewing profiler output¶

There are a few ways to view the information collected by the profiler. You can view it in the console, you can view a more graphical version in a browser, or you can use a vendor tool such as Intel VTune or Nvidia NVProf to view output. For most scenarios the information you need can be obtained with MXNet’s built in profiler support, but if you want to investigate the performance of operators alongside extra context about your hardware (e.g. cache hit rates, or CUDA kernel timings) then profiling jointly with vendor tools is recommended.

print(profiler.dumps())

profiler.dump()

NVProf and Visual Profiler¶

NVProf and Visual Profiler are available in CUDA 9 and CUDA 10 toolkits. You can get a timeline view of CUDA kernel executions, and also analyse the profiling results to get automated recommendations. It is useful for profiling end-to-end training but the interface can sometimes become slow and unresponsive.

You can initiate the profiling directly from inside Visual Profiler or from the command line with nvprof which wraps the execution of your Python script. If it’s not on your path already, you can find nvprof inside your CUDA directory. See this discussion post for more details on setup.

$ nvprof -o my_profile.nvvp python my_profiler_script.py

==11588== NVPROF is profiling process 11588, command: python my_profiler_script.py

==11588== Generated result file: /home/user/Development/incubator-mxnet/ci/my_profile.nvvp

We specified an output file called my_profile.nvvp and this will be annotated with NVTX ranges (for MXNet operations) that will be displayed alongside the standard NVProf timeline. This can be very useful when you’re trying to find patterns between operators run by MXNet, and their associated CUDA kernel calls.

You can open this file in Visual Profiler to visualize the results.

At the top of the plot we have CPU tasks such as driver operations, memory copy calls, MXNet engine operator invocations, and imperative MXNet API calls. Below we see the kernels active on the GPU during the same time period.

Zooming in on a backwards convolution operator we can see that it is in fact made up of a number of different GPU kernel calls, including a cuDNN winograd convolution call, and a fast-fourier transform call.

Selecting any of these kernel calls (the winograd convolution call shown here) will get you some interesting GPU performance information such as occupancy rates (vs theoretical), shared memory usage and execution duration.

Nsight Compute¶

Nsight Compute is available in CUDA 10 toolkit, but can be used to profile code running CUDA 9. You don’t get a timeline view, but you get many low level statistics about each individual kernel executed and can compare multiple runs (i.e. create a baseline).

Further reading¶