如何在Python脚本中直接检测PyTorch是否使用GPU？

嘿，我来给你分享几个在Python脚本里直接检查PyTorch是否使用GPU的实用方法，比单独跑nvidia-smi更贴合PyTorch的实际运行场景～

1. 先确认PyTorch能识别到GPU #

第一步得先搞清楚你的PyTorch环境有没有正确关联GPU，用一行代码就能验证：

import torch

# 输出True表示有可用GPU，False则是没识别到
print(torch.cuda.is_available())

要是返回False，那大概率是CUDA驱动没装对、PyTorch版本和CUDA版本不匹配，或者硬件本身没有GPU，得先解决这个基础问题。

2. 查看张量/数据是否在GPU上 #

PyTorch的计算是基于张量的，只有当张量在GPU上时，才会用GPU运算。你可以直接查看张量的设备信息：

# 创建一个测试张量
test_tensor = torch.tensor([1, 2, 3])
# 打印它所在的设备，默认是CPU
print(test_tensor.device)

# 手动把张量移到GPU上
tensor_on_gpu = test_tensor.to("cuda")
# 现在再看，应该显示cuda:0（第一个GPU设备）
print(tensor_on_gpu.device)

训练的时候，记得把所有输入数据都移到GPU上，不然PyTorch还是会用CPU计算。

3. 验证模型是否部署在GPU上 #

光数据在GPU还不够，模型本身也得移到GPU才行。可以通过检查模型参数的设备来确认：

import torch.nn as nn

# 定义一个简单模型
simple_model = nn.Linear(10, 2)
# 把模型移到GPU
simple_model.to("cuda")

# 查看模型参数的设备，应该返回cuda:0
print(next(simple_model.parameters()).device)

只有模型和数据都在同一个GPU设备上，PyTorch才会启动GPU加速计算。

4. 脚本内实时查看GPU使用率 #

如果想在脚本运行过程中直接看到GPU的负载情况，也可以调用nvidia-smi命令，用Python的subprocess模块实现：

import subprocess

def print_gpu_status():
    # 调用nvidia-smi并捕获输出
    gpu_info = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
    print(gpu_info.stdout)

# 随时调用这个函数查看状态
print_gpu_status()

这样不用跳出脚本，就能实时看到GPU的内存占用、使用率这些细节。

总的来说，核心逻辑就是确保模型和输入数据都部署在GPU设备上，PyTorch就会自动利用GPU进行加速计算啦～

内容的提问来源于stack exchange，提问作者vvvvv