### 如何检测 PyTorch 是否使用 GPU 加速 要检测 PyTorch 是否能够使用 GPU 加速，可以通过以下方法实现。这些方法可以帮助确认系统中是否存在可用的 GPU，以及当前运行环境是否已切换到 GPU 模式。 #### 1. 检查是否有可用的 GPU PyTorch 提供了 `torch.cuda.is_available()` 函数来检测系统中是否安装了 CUDA，并且是否可以使用 GPU。如果返回值为 `True`，则表示系统支持 GPU 加速[^3]。 ```python import torch if torch.cuda.is_available(): print("GPU 可用") else: print("GPU 不可用") ``` #### 2. 获取可用 GPU 的数量 通过 `torch.cuda.device_count()` 可以获取系统中可用 GPU 的数量。如果返回值大于 0，则表示存在可用的 GPU[^2]。 ```python gpu_count = torch.cuda.device_count() print(f"可用 GPU 数量: {gpu_count}") ``` #### 3. 设置设备对象 为了明确指定计算在 CPU 或 GPU 上进行，可以创建一个 `torch.device` 对象。如果 GPU 可用，则选择第一个 GPU (`cuda:0`)；否则，默认使用 CPU。 ```python device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"当前使用的设备: {device}") ``` #### 4. 测试张量是否在 GPU 上 可以通过将张量移动到 GPU 上，并检查其设备属性来验证是否成功切换到 GPU 模式[^4]。 ```python # 创建一个张量 tensor = torch.tensor([1.0, 2.0, 3.0]) # 将张量移动到 GPU if torch.cuda.is_available(): tensor = tensor.to('cuda') # 打印张量所在的设备 print(f"张量所在设备: {tensor.device}") ``` #### 5. 验证运算是否在 GPU 上执行 创建两个张量并进行简单运算，确保它们位于 GPU 上。如果张量在 GPU 上，则所有相关运算也会自动在 GPU 上完成[^4]。 ```python if torch.cuda.is_available(): a = torch.rand(3, 3).to('cuda') b = torch.rand(3, 3).to('cuda') c = a + b print(f"运算结果所在设备: {c.device}") else: print("无法测试 GPU 运算，因为没有可用的 GPU") ``` ### 注意事项 - 如果需要手动设置主 GPU，可以使用 `os.environ.setdefault("CUDA_VISIBLE_DEVICES", "X")` 来限制可见的 GPU 设备[^2]。 - 推荐使用 `try_gpu()` 和 `try_all_gpu()` 等函数来灵活选择设备[^1]。