改变Tensorflow PTXAS定位

Question

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

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

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似乎默认使用CUDA11.0文件夹中的ptxas版本。

我怎样才能将tensorflow指向PTXAS的11.1和11.2版本而不是11.0版本？

Answer 1

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

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

这应该会删除这条信息。