web-dev-qa-db-fra.com

TensorFlow: comment enregistrer l'utilisation de la mémoire GPU (VRAM)?

TensorFlow alloue (pré-) toujours toute la mémoire libre (VRAM) sur ma carte graphique, ce qui est correct car je veux que mes simulations s'exécutent aussi vite que possible sur mon poste de travail.

Cependant, je voudrais enregistrer la quantité de mémoire (en somme) que TensorFlow utilise réellement. De plus, ce serait vraiment bien, si je pouvais également enregistrer la quantité de mémoire utilisée par les tenseurs simples.

Ces informations sont importantes pour mesurer et comparer la taille de la mémoire dont les différentes architectures ML/AI ont besoin.

Des conseils?

pythontensorflowcudnnvram
18
daniel451

Mise à jour, peut utiliser les opérations TensorFlow pour interroger l'allocateur:

# 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())

Vous pouvez également obtenir des informations détaillées sur session.run appelant toutes les allocations de mémoire pendant run en regardant RunMetadata. IE quelque chose comme ça

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

Voici un exemple de bout en bout - prenez un vecteur de colonne, un vecteur de ligne et ajoutez-les pour obtenir une matrice d'additions:

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

Si vous ouvrez run.txt vous verrez des messages comme celui-ci:

  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

Donc, ici, vous pouvez voir que a et b ont alloué 52 octets chacun (13 * 4), et le résultat a alloué 676 octets.

18
Yaroslav Bulatov