TensorFlow检测GPU是否可用

背景

安装了TensorFlow,训练的时候死活用不到GPU,特别是在docker环境中,GPU是否可检查更为重要

依赖条件

已经安装好python、响应版本GPU的驱动和nvidia的驱动

检测代码

import tensorflow as tf

print(tf.__version__)
print(tf.reduce_sum(tf.random.normal([1000, 1000])))
print(tf.config.list_physical_devices('GPU'))

输出结果可以查看是否有GPU信息

torch的GPU检测

import torch

# 检查CUDA是否可用
if torch.cuda.is_available():
    # 输出可用的CUDA设备数量
    print(torch.cuda.device_count(), "GPU(s) available.")
    
    # 输出当前GPU设备的名称
    print("Current GPU:", torch.cuda.get_device_name(0))
else:
    print("CUDA is not available. Using CPU.")