关于Tensorflow与Gym MuJoCo同时使用的问题

引子

        大家好，今天我们稍微介绍一下如何处理TensorFlow与Gym MuJoCo同时使用将会出错的问题。

问题

        在同时使用TensorFlow-GPU和Gym MuJoCo时，可能会报错：“failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED”、“Blas GEMM launch failed”。这是因为我们的TensorFlow-GPU总是倾向于占用全部GPU，而在我们对MuJoCo中的环境调用render()时，也需要用到GPU，所以引发了上述问题。

解决办法

        TensorFlow为我们提供了GPU占用方式的控制接口，一般用tf.ConfigProto()对tf.Session()进行参数配置：

sess = tf.Sesssion(config=tf.ConfigProto(...), ...)

        首先我们来了解一下tf.ConfigProto()。先介绍两个常用的参数：

• log_device_placement=True：是否打印设备分配日志；
• allow_soft_placement=True：若我们指定的设备不存在，则允许TF自动分配设备。

        下面介绍如何控制GPU资源的分配。

allow_growth

        这一选项将允许我们为TensorFlow一开始分配较少的GPU，然后按需慢慢增加，不过因为不会释放内存，所以会导致GPU碎片。

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Sesstion(config=config)

per_process_gpu_memory_fraction

        设置TensorFlow在每个GPU上的份额占比，0.5表示50%。

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = True
sess = tf.Sesstion(config=config)

GPU选择

        我们可以在设置环境变量CUDA_VISIBLE_DEVICES=0，表示使用GPU 0，设置CUDA_VISIBLE_DEVICES=0, 1表示使用GPU 0, 1：

$ export CUDA_VISIBLE_DEVICES=0,1

        或者在python脚本中：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

问题解决

        我的电脑只有一块GPU，所以尝试不了第三种方法，即在不同的GPU上运行TensorFlow和MuJoCo render。不过第一种方法allow_growth和第二种方法per_process_gpu_memory_fraction都是可以解决上面的问题的。