更改 Tensorflow PTXAS 位置

Question

我目前正在尝试使用 Tensorflow 2.4.0 和运行 CUDA 11.0 和 CUDNN 8 的 RTX 3070 自定义训练神经网络。

我遇到了这个奇怪的问题，我可以训练模型，但实际上无法获得任何输出，因为当我运行时：

output = model(x) 我收到以下消息，并且我的 jupyter 内核自动重新启动。

2021-01-08 20:52:53.437668: W tensorflow/stream_executor/gpu/asm_compiler.cc:191] Falling back to the CUDA driver for PTX compilation; ptxas does not support CC 8.6
2021-01-08 20:52:53.437690: W tensorflow/stream_executor/gpu/asm_compiler.cc:194] Used ptxas at /usr/local/cuda-11.0/bin/ptxas
2021-01-08 20:52:53.438427: W tensorflow/stream_executor/gpu/redzone_allocator.cc:314] Unimplemented: /usr/local/cuda-11.0/bin/ptxas ptxas too old. Falling back to the driver to compile.
Relying on driver to perform ptx compilation. 
Modify $PATH to customize ptxas location.

作为测试，我安装了 CUDA 11.1 和 11.2 并相应地重新调整了 $PATH 变量，但tensorflow似乎默认使用 CUDA 11.0 文件夹中的 ptxas 版本。

我该怎么做才能将tensorflow指向PTXAS 11.1和11.2版本而不是11.0版本？

Answer 1

将 nVidia 工具包 bin 目录添加到您的路径中，例如：

export PATH=/depot/cuda/cuda-11.2/bin:$PATH

这应该删除该消息。