为什么在 Stable Baselines 3 中使用 GPU 比使用 cpu 慢？

Question

当使用 PPO 使用 Stable Baselines 3 训练“CartPole”环境时，我发现使用 cuda GPU 训练模型的速度几乎是仅使用 cpu 训练模型的两倍（无论是在 google colab 中还是在本地）。

我认为使用 cuda 来完成这种任务会更快，我做错了什么吗？

我的代码是：

import gym
import time
from stable_baselines3 import PPO

env = gym.make("CartPole-v1")

t1 = time.time()
model = PPO("MlpPolicy", env, verbose=0, device="cuda")
model.learn(total_timesteps=10_000)
print(f"Time with cuda : {time.time()-t1:.2f}s")

t1 = time.time()
model = PPO("MlpPolicy", env, verbose=0, device="cpu")
model.learn(total_timesteps=10_000)
print(f"Time with cpu : {time.time()-t1:.2f}s")

env.close()

输出为：

Time with cuda : 21.76s
Time with cpu : 13.33s

Answer 1

正如https://stable-baselines3.readthedocs.io/en/master/modules/ppo.html中提到的，StableBaseline3 的 PPO 应该在 CPU 上运行。

StableBaselines3 PPO 注释