### 如何在 PyTorch 中检测当前运行环境以及 GPU 设备 在 PyTorch 中，可以通过一系列方法来检测当前运行环境是否使用了 GPU 或其他特定设备。以下是详细的实现方式： #### 检测 GPU 是否可用 可以使用 `torch.cuda.is_available()` 方法来判断是否有可用的 GPU。 ```python import torch # 判断 GPU 是否可用 is_cuda_available = torch.cuda.is_available() print(f"CUDA 可用: {is_cuda_available}") ``` 如果返回值为 `True`，则表示系统中有可用的 GPU 并已成功配置 CUDA[^1]。 #### 获取 GPU 数量 当有多个 GPU 存在于环境中时，可通过以下代码获取具体的 GPU 数量： ```python if is_cuda_available: gpu_count = torch.cuda.device_count() print(f"GPU 数量: {gpu_count}") else: print("没有可用的 GPU。") ``` #### 查看具体 GPU 信息 对于每一块可用的 GPU，可进一步查询其名称和其他属性： ```python for i in range(torch.cuda.device_count()): gpu_name = torch.cuda.get_device_name(i) print(f"GPU {i} 名称: {gpu_name}") ``` #### 当前默认使用的 GPU 通过以下代码可以获得当前正在使用的 GPU 索引及其名称： ```python if is_cuda_available: current_device_index = torch.cuda.current_device() print(f"当前使用的 GPU 索引: {current_device_index}") current_gpu_name = torch.cuda.get_device_name(current_device_index) print(f"当前 GPU 名称: {current_gpu_name}") ``` #### 测试张量是否位于 GPU 上 为了验证某个张量是否已经加载到 GPU 上，可以创建一个简单的张量并打印其设备位置： ```python tensor_on_cpu = torch.tensor([0.0, 1.0]) print(f"CPU 张量所在设备: {tensor_on_cpu.device}") # 输出应为 'cpu' if is_cuda_available: tensor_on_gpu = tensor_on_cpu.to('cuda') print(f"GPU 张量所在设备: {tensor_on_gpu.device}") # 输出应类似于 'cuda:0' ``` #### 高级功能：检查 Tensor Core 和 BF16 支持情况 某些高级特性可能依赖于硬件的支持程度，例如 Tensor Cores 和 BFloat16 数据格式。这些也可以被轻松检测： ```python supports_bf16 = torch.cuda.is_bf16_supported() print(f"是否支持 BF16 数字格式: {'支持' if supports_bf16 else '不支持'}") has_tensor_core_support = torch.cuda.get_device_capability()[0] >= 7 print(f"是否支持 Tensor Core: {'支持' if has_tensor_core_support else '不支持'}") ``` 以上代码片段展示了如何全面地评估当前系统的 GPU 运行环境，并确认 PyTorch 能否充分利用硬件资源[^3]。 ---