nvtx-plugins 0.1.8

先决条件

• Linux

• Python 3.4+

• NVIDIA GPU + CUDA工具包10.0或更高版本

• TensorFlow 1.13或更高版本

安装NVTX-Plugins

该软件包可以从PyPI安装

# Stable release
pip install nvtx-plugins

# Pre-release (may present bugs)
pip install nvtx-plugins --pre

该软件包还可在GitHub上下载：https://github.com/NVIDIA/nvtx-plugins/releases

pip install nvtx-plugins*.tar.gz

从源代码安装

您可以从源代码构建和安装该软件包

python setup.py sdist
pip install dist/nvtx-plugins*.tar.gz

对于开发目标，您可以使用以下命令直接从源代码安装软件包：

python setup.py install

我们建议在NVIDIA的NGC TensorFlow容器内构建软件包：https://ngc.nvidia.com/catalog/containers/nvidia:tensorflow

有关如何开始使用NGC容器的更多信息，请参阅以下NVIDIA GPU云文档和Deep Learning DGX文档的部分：从NVIDIA GPU云开始使用、访问并从NGC容器注册表中拉取以及运行TensorFlow。

构建文档

通过运行以下命令构建文档：

cd docs
pip install -r requirements.txt
make html

文档文件将生成在docs/build/html

构建文档不需要安装NVTX插件。不过，由于Sphinx中的问题，**只支持Python 3.7构建文档**。

向图中添加标记

通过将计算图的部分包装在开始和结束操作中添加标记。这些操作是恒等操作（不修改地将输入传递到输出）但它们有生成nvtx标记的副作用。

import nvtx.plugins.tf as nvtx_tf

x, nvtx_context = nvtx_tf.ops.start(x, message='Dense 1-3',
    domain_name='Forward', grad_domain_name='Gradient')
x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_1')
x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_2')
x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_3')
x = nvtx_tf.ops.end(x, nvtx_context)
x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_4')

为了方便，该软件包还提供了一个函数装饰器

@nvtx_tf.ops.trace(message='Dense Block', domain_name='Forward',
                   grad_domain_name='Gradient')
def dense_block(x):
    x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_1')
    x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_2')
    x = tf.layers.dense(x, 1000, activation=tf.nn.relu, name='dense_3')
    return x

更详细的示例可以在examples/中找到，还可以查看文档以获取有关其他工作流程的更多信息，包括会话钩子、Keras层和回调。

可视化范围

NVTX需要一个记录器来注册生成的事件和范围，我们将使用NVIDIA Nsight Systems来捕获这些事件。

使用nsys（预安装在NVIDIA的NGC TensorFlow容器中）运行您的代码以生成一个qdrep文件

nsys profile -d 60 \
    -w true \
    --sample=cpu \
    -t 'nvtx,cuda' \
    -o ./generated_timeline \
    python ./network.py

生成的qdrep可以使用Nsight Systems查看。

Nsight Systems和nsys也可以从NVIDIA开发者网站下载。

有关nsys和Nsight Systems的更多详细信息，请参阅此处。