TensorFlow の推論処理で複数のGPUを使用する (C++)
TensorFlow の推論処理を特に指定なく実行すると GPU 0 だけが使用されます。これを変更する方法をいくつか紹介しようと思います。
環境変数
環境変数 CUDA_VISIBLE_DEVICES で使用するGPUを指定することができます。
例
set CUDA_VISIBLE_DEVICES=0
set_visible_device_list
以下のように set_visible_device_list でGPUを指定することが可能です。引数で指定する値は上記の環境変数 CUDA_VISIBLE_DEVICES で指定するものと同じです。
例
tensorflow::Session::Options options;
options.config.mutable_gpu_options()->set_visible_device_list("0");
ただしこの場合、複数のスレッドで異なる値を指定した場合に以下のエラーが発生してしまいます。
Already exists: TensorFlow device (GPU:0) is being mapped to multiple CUDA devices (0 now, and 2 previously), which is not supported. This may be the result of providing different GPU configurations (ConfigProto.gpu_options, for example different visible_device_list) when creating multiple Sessions in the same process. This is not currently supported, see https://github.com/tensorflow/tensorflow/issues/19083
SetDefaultDevice
SetDefaultDevice で使用するデバイスを指定することができます。この場合、グラフに GPU でサポートされていないノードが存在することがあるので、これを回避するために set_allow_soft_placement で true を指定しておきます。
tensorflow::Session::Options options;
options.config.set_allow_soft_placement(true);
tensorflow::GraphDef graph_def;
// GraphDef の読み込みは省略します
::tensorflow::graph::SetDefaultDevice("/gpu:0", &graph_def);
その他
以下の処理は効果がありませんでした。
以下の設定が可能との情報も見たのですが、少なくとも私の環境では無効でした。
tensorflow::Session::Options options;
options.config.mutable_device_count()->insert({"GPU", 2});