TensorFlow: cách đăng nhập sử dụng bộ nhớ GPU (VRAM)?

Question

TensorFlow luôn luôn (trước) phân bổ tất cả bộ nhớ miễn phí (VRAM) trên card đồ họa của tôi, điều này là ổn vì tôi muốn mô phỏng của mình chạy nhanh nhất có thể trên máy trạm của tôi.

Tuy nhiên, tôi muốn ghi lại số lượng bộ nhớ (tính bằng tổng) TensorFlow thực sự sử dụng. Ngoài ra nó sẽ được thực sự tốt đẹp, nếu tôi cũng có thể đăng nhập bao nhiêu bộ nhớ duy nhất tensors sử dụng.

Thông tin này rất quan trọng để đo lường và so sánh kích thước bộ nhớ mà các kiến ​​trúc ML/AI khác nhau cần.

Bất kỳ mẹo nào?

Answer 1

Update, có thể sử dụng ops TensorFlow để truy vấn cấp phát:

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

Ngoài ra bạn có thể nhận được thông tin chi tiết về session.run cuộc gọi bao gồm tất cả các bộ nhớ được phân bổ trong run cuộc gọi bằng cách nhìn vào RunMetadata. IE một cái gì đó như thế này

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

Dưới đây là một ví dụ end-to-end - mất vector cột, vector hàng và thêm chúng để có được một ma trận bổ sung:

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

Nếu bạn mở run.txt bạn' ll xem các tin nhắn như thế này:

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 

Vì vậy, ở đây bạn có thể thấy rằng a và b phân bổ mỗi 52 byte (13 * 4), và kết quả phân bổ 676 byte.