[原创] 强化学习框架 rlpyt 使用GPU训练模型

rlpyt 是BAIR(Berkeley Artificial Intelligence Research,伯克利人工智能研究所)开源的一个强化学习(RL)框架。我之前写了一篇它的简介。 

rlpyt 可利用GPU来训练模型,理论上速度会比CPU快。

安装GPU版的PyTorch
在Anaconda下,利用rlpyt自带的 linux_cuda10.yml 或 linux_cuda9.yml 来安装 rlpyt 依赖的Anaconda env(取决于你使用的CUDA Version是多少)。
使用 nvidia-smi 命令查看 CUDA Version:

NVIDIA-SMI 410.72       Driver Version: 410.72       CUDA Version: 10.0
我的是 10.0,因此修改 linux_cuda10.yml,把下面这行:
cudatoolkit=10.
改成:
cudatoolkit=10.0
我测试下来,如果不改,rlpyt env也可以安装成功,但是装上的是CPU版的PyTorch;不仅如此,如果你服务器上的CUDA Version是10.1的,你也要把 yml 配置文件里改成 10.0,否则装上的也是CPU版的PyTorch:

~$ conda list | grep pytorch
_pytorch_select           0.1                       cpu_0    defaults
pytorch                   1.2.0           cpu_py37h00be3c6_0    defaults

文章来源:https://www.codelast.com/
改了之后装上的就是GPU版的PyTorch了:

~$ conda list | grep pytorch
_pytorch_select           0.2                       gpu_0    defaults
pytorch                   1.2.0           cuda100py37h938c94c_0    defaults

 在模型训练程序调度GPU
具体应该怎么写GPU相关的程序,取决于你要实现什么样的功能。例如,使用同步/异步模式,它们对应的Sampler是不同的,等等。
涉及到具体的代码,这里难以给出一个详细的说明,可以参考 rlpyt 的examples代码。
文章来源:https://www.codelast.com/
 确认真的在使用GPU训练
依然使用 nvidia-smi 命令来检测,正常情况下,在模型开始训练后,会看到 Processes 列表中有你的模型训练进程:
3     17754      C   .../.anaconda/envs/rlpyt/bin/python   459MiB
此外,在模型训练的过程中,查看GPU使用率应该不为0:
dstat --nvidia-gpu -af
这样就OK了。
据我实测,在我的应用场景中,同样功能的强化学习程序,GPU版比CPU版的模型训练速度快50%左右,GPU单卡使用率最高约为14%(只使用了一张卡,Nvidia P40显卡)。
文章来源:https://www.codelast.com/
➤➤ 版权声明 ➤➤ 
转载需注明出处:codelast.com 
感谢关注我的微信公众号(微信扫一扫):

wechat qrcode of codelast

发表评论

电子邮件地址不会被公开。 必填项已用*标注