web-dev-qa-db-ja.com

TensorFlow：GPUメモリ（VRAM）使用率をログに記録する方法？

TensorFlowは常にグラフィックカードにすべての空きメモリ（VRAM）を（事前に）割り当てます。これは、私のワークステーションでシミュレーションを可能な限り高速に実行したいためです。

ただし、TensorFlowが実際に使用しているメモリ（合計）の量を記録したいと思います。さらに、シングルテンソルが使用するメモリの量もログに記録できれば、それは本当にすばらしいでしょう。

この情報は、さまざまなML/AIアーキテクチャに必要なメモリサイズを測定して比較するために重要です。

任意のヒント？

pythontensorflowcudnnvram
18
daniel451

更新。TensorFlow演算を使用してアロケータをクエリできます。

# maximum across all sessions and .run calls so far
sess.run(tf.contrib.memory_stats.MaxBytesInUse())
# current usage
sess.run(tf.contrib.memory_stats.BytesInUse())

また、runを参照することにより、RunMetadata呼び出し中に割り当てられているすべてのメモリを含むsession.run呼び出しに関する詳細情報を取得できます。 IEこのようなもの

run_metadata = tf.RunMetadata()
sess.run(c, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True), run_metadata=run_metadata)

これがエンドツーエンドの例です-列ベクトル、行ベクトルを取り、それらを追加して追加の行列を取得します。

import tensorflow as tf

no_opt = tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0,
                             do_common_subexpression_elimination=False,
                             do_function_inlining=False,
                             do_constant_folding=False)
config = tf.ConfigProto(graph_options=tf.GraphOptions(optimizer_options=no_opt),
                        log_device_placement=True, allow_soft_placement=False,
                        device_count={"CPU": 3},
                        inter_op_parallelism_threads=3,
                        intra_op_parallelism_threads=1)
sess = tf.Session(config=config)

with tf.device("cpu:0"):
    a = tf.ones((13, 1))
with tf.device("cpu:1"):
    b = tf.ones((1, 13))
with tf.device("cpu:2"):
    c = a+b

sess = tf.Session(config=config)
run_metadata = tf.RunMetadata()
sess.run(c, options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE, output_partition_graphs=True), run_metadata=run_metadata)
with open("/tmp/run2.txt", "w") as out:
  out.write(str(run_metadata))

run.txtを開くと、次のようなメッセージが表示されます。

  node_name: "ones"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322108320
      }
  ....

  node_name: "ones_1"

      allocation_description {
        requested_bytes: 52
        allocator_name: "cpu"
        ptr: 4322092992
      }
  ...
  node_name: "add"
      allocation_description {
        requested_bytes: 676
        allocator_name: "cpu"
        ptr: 4492163840

したがって、ここでは、aとbがそれぞれ52バイト（13 * 4）を割り当て、結果が676バイトを割り当てていることがわかります。

18
Yaroslav Bulatov