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});

参考

• TensorflowでGPUを制限・無効化する
• Select specific gpu for the session of tensorflow c++ api

tensorflowの関数は説明してもわかりにくいので、使用例のみで解説します。

tf.Variable()

import tensorflow as tf
x = tf.Variable(10 , name="x")
y = x * 5
print(y)

出力
Tensor("mul_1:0", shape=(), dtype=int32)

上記のコードを見ると、出力が50のように思えますが、テンソルが出力されます。
では、次に、50が出力されるにはどのようにすればよいでしょうか。

x = tf.Variable(10 , name="x")
y = x * 5
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
sess.run(y)

出力50
tensorflowでは、tf.sessionを通じて計算グラフを構築して実行しなければいけません。tf.global_variables_initializer()は、計算グラフ内の変数の初期化の役割を果たします

tf.placeholder()

tf.placeholder(dtype, shape=None, name=None)

dtype: 変数に代入される型を代入します。
shape: 割り当てる変数のshape
name: 変数の名前

返り値
直接評価されない、feedされる変数のTensor

プログラム例

x = tf.placeholder(tf.int32)
y = tf.constant(1)
z = x + y
with tf.Session() as sess:
  print(sess.run(z, feed_dict={x: 2}))

出力
3
x=2を実行時に定義
次に、、tf.placeholderに配列を渡します。

x = tf.placeholder(tf.int32, shape=[2])
y = tf.placeholder(tf.int32, shape=[2])
z = x*y
with tf.Session() as sess:
  print(sess.run(z, feed_dict={x: [2, 1], y: [1, 2]}))

出力
[2 2]

tf.placeholder_with_default

tf.placeholder_with_defaultを使うことで、初期値を決められる、かつ、実行時に変更もできます。

x = tf.placeholder_with_default(1, shape=[])
y = tf.constant(1)
with tf.Session() as sess:
  print(sess.run(x + y))

出力
0

x = tf.placeholder_with_default(1, shape=[])
y = tf.constant(1)
with tf.Session() as sess:
  print(sess.run(x + y, feed_dict={x: -1}))

出力
2

tf.shape

tf.shapeは動的に変更されうるshapeに使用します。

x = tf.placeholder(tf.int32, shape=[None, 3])
size = tf.shape(x)[0]
sess = tf.Session()
sess.run(size, feed_dict={x: [[1, 2, 3], [1, 2, 3]]})

出力
2

.get_shape

.get_shapeは変更されないshapeに使用します。

x = tf.constant(1, shape=[2, 3, 4])
x.get_shape()[1]

出力
Dimension(3)

参考文献
Tensorflowで学ぶディープラーニング入門