【分散学習】 TensorflowにおけるCPU/GPU使い分け

TensorflowにおけるプロセスごとのCPU/GPU使い分けについて記述します。

背景

Ape-x,DISTRIBUTEDPRIORITIZEDEXPERIENCEREPLAY
R2D2,Recurrent Experience Replay in Distributed Reinforcement Learning

強化学習において、Ape-x、R2D2のように経験の獲得はCPUを用いて並列化し、その経験を中央のGPUが学習する手法が提案されています。
CPUは並列化しやすいため、学習効率の改善が期待できます。

環境

Ubuntu18.04
Python 3.6.8
Tensorflow 1.12.0
CUDA 10.1

CPU/GPU使い分け

解決策

Ape-x,R2D2ではマルチプロセスを用い、GPU計算をするLearnerと、CPU計算をするActorを生成します。
このとき、プロセス間のCPUとGPUの使い分けは下記コードで実現出来ます。

import os,multiprocessing

def hoge_cpu():
    #CPU計算をしたい
    os.environ["CUDA_VISIBLE_DEVICES"] = "" #""には何も書かない
    ...

def hoge_gpu():
    #GPU計算をしたい
    ...

multiprocessing.Process(target=hoge)
multiprocessing.Process(target=hoge_gpu)

os.environ["CUDA_VISIBLE_DEVICES"]がポイントです。これを書くだけです。

別の実装

他の方の実装では下記のコードがありました。

import tensorflow as tf
with tf.device("/gpu:0"):
    multiprocessing.Process(target=hoge_gpu)
with tf.device("/cpu:0"):
    multiprocessing.Process(target=hoge_cpu)

しかしこの実装ではCPU側の実行にGPUのメモリ領域が確保されていました。プロセス数に余裕があってもGPUのメモリサイズが限界となり、並列化の数が小さくなってしまいます。