更改 Tensorflow PTXAS 位置

更改 Tensorflow PTXAS 位置

我目前正在尝试使用 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 版本?

答案1

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

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

这应该会删除该消息。

相关内容