在没有CUDA_VISIBLE的情况下在keras后端中指定GPU

Question

我有一个问题，我必须在同一程序中同时包含keras和pytorch框架。冻结的特征提取器位于Keras中。当我在多GPU模式（4 GPU）中运行Keras特征提取器时，它占用了所有GPU的执行时间，而下半部分（pytorch模型）没有空间，我想为keras模型分配2 gpu，将其他gpu分配给pytorch模型。如果我设置了cuda.visible_devices='0,1'，则pytorch也不允许同时使用其他3个GPU，这是一个棘手的情况。 另一个解决方法是设置keras后端会话：

config = tf.ConfigProto( device_count = {'GPU': 1} )
sess = tf.Session(config=config)
keras.backend.set_session(sess)

但是在这里，我想指定会话可以使用哪些GPU（按ID）。 我遇到的错误是在下半部分（pytorch部分）之前将数据推送到gpu时，CUDA内存不足错误，因为keras已经声明了gpu占用。

RuntimeError: CUDA out of memory. Tried to allocate 204.00 MiB (GPU 2; 10.92 GiB total capacity; 14.54 MiB already allocated; 151.50 MiB free; 9.46 MiB cached)

我可以使用pytorch指定GPU。我只想对keras做同样的事情。