如何使用 os.environ['CUDA_VISIBLE_DEVICES'] ="" 将 GPU 计数设置为 0？答案

【问题标题】：How to set GPU count to 0 using os.environ['CUDA_VISIBLE_DEVICES'] =""?如何使用 os.environ['CUDA_VISIBLE_DEVICES'] ="" 将 GPU 计数设置为 0？
【发布时间】：2021-04-29 05:32:53
【问题描述】：

所以我在我的系统中配置了以下 GPU：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 461.33       Driver Version: 461.33       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100S-PCI...  TCC  | 00000000:3B:00.0 Off |                    0 |
| N/A   30C    P0    25W / 250W |      1MiB / 32642MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100S-PCI...  TCC  | 00000000:D8:00.0 Off |                    0 |
| N/A   31C    P0    25W / 250W |      1MiB / 32642MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

现在，我必须通过 python 设置环境，例如 GPU 计数 = 0。在从各种来源学习后，我尝试了以下方法：

import os
os.environ["CUDA_VISIBLE_DEVICES"]=""
import torch
torch.device_count()

但是，它仍然给我输出“2”，就像系统中的 2 个 GPU 一样。如何设置环境，使其输出“0”？

任何其他方式，将计数设置为“0”也是值得赞赏的，但它应该与任何 ML-Library 无关。（例如，我不能使用device = torch.device("cpu")，因为这仅适用于 Pytorch 而不适用于其他库）

【问题讨论】：

您是否尝试将此值全局设置为 0 或让所有库在所有实例中忽略 CUDA？我认为torch.cuda.device_count() 最终会直接与 CUDA API 本身对话，因此将始终返回实际计数，而不是环境变量设置的值。
是的，我希望所有 ml 库中的所有实例都认为“没有 GPU 可用/可见”
无法重现。但我必须使用torch.cuda.device_count()，而不是torch.device_count()。它正确地将 GPU 的数量更改为 0。

标签： python-3.x pytorch nvidia python-os

【解决方案1】：

为防止您的 GPU 被使用，请设置os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

【讨论】：

【解决方案2】：

最简单的方法是使用正确的环境集运行python。例如，在 Linux 上

CUDA_VISIBLE_DEVICES="" python ...

以下也应该有效：

os.environ["CUDA_VISIBLE_DEVICES"]=""

但这必须在您首先导入torch之前完成。

我认为你的情况是你必须更早地导入torch，也许是通过一些使用torch的库间接导入的。

【讨论】：

【解决方案3】：

os.environ["CUDA_VISIBLE_DEVICES"]="0,1"
torch.cuda.device_count() # 结果是 2

os.environ["CUDA_VISIBLE_DEVICES"]="0"
torch.cuda.device_count() # 结果为 1，使用第一个 GPU

os.environ["CUDA_VISIBLE_DEVICES"]="1"
torch.cuda.device_count() # 结果为 1，使用第二个 GPU

【讨论】：

感谢您的回答。然而，这个答案并不能解决手头的问题。如果您知道如何将torch.cuda.device_count() 设为零，请考虑更新答案。