在 Tensorflow 2.0 中检查 python 使用的 GPU 内存

Question

在 SO 上有几个线程 here 和 here 涵盖了如何使用 conrib 库和会话在 python 中获取 Tensorflow 使用的 GPU 内存，但是我们如何在 TF 2.0 中在急切执行中做到这一点（ contrib 库不适用于 2.0)？

Answer 1

目前看来，这个选项在 TF 2 中不可用。一些替代方案包括：

• 按照in this issue 的说明，为 NVIDIA 管理库使用 python 绑定
• 通过nvidia-smi 命令获取信息

对于第二个选项，您可以执行类似于this answer 的操作来获取某些 GPU 中当前使用的内存。

我们首先得到 gpu 的初始状态，然后我们设置 TF 不使用超过所需的内存（默认是使用所有可用内存），然后我们得到 gpu 的当前状态。

import subprocess as sp
import tensorflow as tf

def gpu_memory_usage(gpu_id):
    command = f"nvidia-smi --id={gpu_id} --query-gpu=memory.used --format=csv"
    output_cmd = sp.check_output(command.split())
    
    memory_used = output_cmd.decode("ascii").split("\n")[1]
    # Get only the memory part as the result comes as '10 MiB'
    memory_used = int(memory_used.split()[0])

    return memory_used

# The gpu you want to check
gpu_id = 0

initial_memory_usage = gpu_memory_usage(gpu_id)

# Set up the gpu specified
gpu_physical_devices = tf.config.list_physical_devices('GPU')
for device in gpu_physical_devices:
    if int(device.name.split(":")[-1]) == gpu_id:
        device_to_be_used = device
        # Set memory growth for TF to not use all available memory of the GPU
        tf.config.experimental.set_memory_growth(device, True)

# Just to be sure that we are only using the required gpu
tf.config.set_visible_devices([device_to_be_used], 'GPU')


# Create your model here
# Do cool stuff ....

latest_gpu_memory = gpu_memory_usage(gpu_id)
print(f"(GPU) Memory used: {latest_gpu_memory - initial_memory_usage} MiB")

请注意，我们在这里做了一些假设，例如，没有其他进程与我们的进程同时启动，并且已经在 GPU 中运行的其他进程不需要使用更多内存。