Stable Baselines 3 DQN模型识别GPU却拒绝使用CUDA的问题

我在Stable Baselines 3的自定义环境里碰到了DQN模型拒绝使用GPU的问题。

初始化模型的代码如下：

model = DQN("MlpPolicy", env, device="cuda")

相关软硬件信息 #

• 显卡：RTX 2070 Super
• 已安装CUDA版本：10.1
• 已安装cudNN版本：7.5.0（适配CUDA 10.1）
• 已安装Tensorflow版本：2.2

单独用Tensorflow做机器学习时CUDA能正常工作，但在Stable Baselines 3里始终默认用CPU。不过在创建模型前打印可用CUDA设备时，显示有一个可用设备（我的RTX 2070 Super）。

Tensorflow初始化时所有插件都加载正常，没有报错。我试过更换多个CUDA版本、Tensorflow版本，也装过特定的tensorflow-gpu包，但都没解决问题，也没找到相关解决方案。

运行日志 #

2023-02-01 16:28:20.409397: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cudart64_101.dll
2023-02-01 16:28:21.691033: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library nvcuda.dll
2023-02-01 16:28:21.713878: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1561] Found device 0 with properties: 
pciBusID: 0000:2d:00.0 name: NVIDIA GeForce RTX 2070 SUPER computeCapability: 7.5
coreClock: 1.815GHz coreCount: 40 deviceMemorySize: 8.00GiB deviceMemoryBandwidth: 417.29GiB/s
2023-02-01 16:28:21.714097: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cudart64_101.dll
Num GPUs Available:  1
Using cpu device
Wrapping the env with a `Monitor` wrapper
Wrapping the env in a DummyVecEnv.
2023-02-01 16:28:21.717235: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cublas64_10.dll
2023-02-01 16:28:21.719513: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cufft64_10.dll
2023-02-01 16:28:21.720370: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library curand64_10.dll
2023-02-01 16:28:21.723773: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cusolver64_10.dll
2023-02-01 16:28:21.725979: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cusparse64_10.dll
2023-02-01 16:28:21.732353: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cudnn64_7.dll
2023-02-01 16:28:21.732490: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1703] Adding visible gpu devices: 0

内容的提问来源于stack exchange，提问作者yuzuki