查看、指定使用的 GPU 数量和编号

JjWang__HUST2024-07-27 8:20

在使用 PyTorch 框架时，可以通过以下步骤查看可用的 GPU 数量，指定使用的 GPU 编号，并在代码中体现这一点。下面以2个GPU为例：

目录

  • 一、脚本代码块实现
    • [1. 查看可用的 GPU](#1. 查看可用的 GPU)
    • [2. 指定使用 GPU 的数量和编号](#2. 指定使用 GPU 的数量和编号)
      • [使用 `CUDA_VISIBLE_DEVICES` 环境变量](#使用 CUDA_VISIBLE_DEVICES 环境变量)
      • [将模型和数据移动到指定的 GPU](#将模型和数据移动到指定的 GPU)
      • [使用 `DataParallel` 进行多 GPU 并行计算](#使用 DataParallel 进行多 GPU 并行计算)
    • [3. 示例完整代码](#3. 示例完整代码)
  • 二、命令行实现
    • [1. 使用 `CUDA_VISIBLE_DEVICES` 环境变量](#1. 使用 CUDA_VISIBLE_DEVICES 环境变量)
      • [只使用 GPU 0](#只使用 GPU 0)
      • [使用 GPU 0 和 GPU 1](#使用 GPU 0 和 GPU 1)
    • [2. 在 Python 代码中自动检测可用 GPU](#2. 在 Python 代码中自动检测可用 GPU)
    • [3. 结合 `CUDA_VISIBLE_DEVICES` 和 PyTorch 代码](#3. 结合 CUDA_VISIBLE_DEVICES 和 PyTorch 代码)

一、脚本代码块实现

1. 查看可用的 GPU

python 复制代码 
import torch

if torch.cuda.is_available():
    num_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {num_gpus}")
    for i in range(num_gpus):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
    print("No GPUs available")

2. 指定使用 GPU 的数量和编号

在 PyTorch 中，可以通过以下几种方法来指定使用的 GPU。

使用 CUDA_VISIBLE_DEVICES 环境变量

可以在代码中设置环境变量来指定使用的 GPU：

python 复制代码 
import os

# 只使用 GPU 0 和 GPU 1
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

import torch

# 检查可用的 GPU
if torch.cuda.is_available():
    num_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {num_gpus}")
    for i in range(num_gpus):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
    print("No GPUs available")

将模型和数据移动到指定的 GPU

可以将模型和数据移动到指定的 GPU。例如：

python 复制代码 
import torch

# 检查是否有 GPU 可用
if torch.cuda.is_available():
    # 将模型移动到 GPU 0
    device0 = torch.device('cuda:0')
    model = MyModel().to(device0)
    
    # 将数据移动到 GPU 1
    device1 = torch.device('cuda:1')
    data = data.to(device1)
else:
    device = torch.device('cpu')
    model = MyModel().to(device)
    data = data.to(device)

使用 DataParallel 进行多 GPU 并行计算

可以使用 torch.nn.DataParallel 来让模型在多个 GPU 上并行运行：

python 复制代码 
import torch
import torch.nn as nn

# 假设模型已经定义好
model = MyModel()

if torch.cuda.is_available():
    # 使用 DataParallel 并行运行在多个 GPU 上
    model = nn.DataParallel(model, device_ids=[0, 1])
    model = model.cuda()
else:
    device = torch.device('cpu')
    model = model.to(device)

3. 示例完整代码

以下是一个完整的示例，演示了如何在 PyTorch 中查看 GPU 数量、指定使用 GPU，并在代码中体现：

python 复制代码 
import os
import torch
import torch.nn as nn

# 设置使用的 GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

# 查看可用的 GPU
if torch.cuda.is_available():
    num_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {num_gpus}")
    for i in range(num_gpus):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
else:
    print("No GPUs available")

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer = nn.Linear(10, 10)

    def forward(self, x):
        return self.layer(x)

model = MyModel()

# 将模型移动到 GPU
if torch.cuda.is_available():
    model = nn.DataParallel(model, device_ids=[0, 1])
    model = model.cuda()

# 示例数据
data = torch.randn(5, 10).cuda()

# 前向传播
output = model(data)
print(output)

在命令行中，可以使用环境变量 CUDA_VISIBLE_DEVICES 来指定使用的 GPU 数量和编号。这种方法对 PyTorch 非常有效。以下是详细步骤：

二、命令行实现

1. 使用 CUDA_VISIBLE_DEVICES 环境变量

在运行 Python 脚本时，通过设置 CUDA_VISIBLE_DEVICES 环境变量来指定哪些 GPU 可见。例如：

只使用 GPU 0

bash 复制代码 
CUDA_VISIBLE_DEVICES=0 python main.py

使用 GPU 0 和 GPU 1

bash 复制代码 
CUDA_VISIBLE_DEVICES=0,1 python main.py

2. 在 Python 代码中自动检测可用 GPU

在 Python 代码中，可以自动检测可用的 GPU，并根据需要配置模型和数据。以下是示例代码：

python 复制代码 
import torch

def print_available_gpus():
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        print(f"Available GPUs: {num_gpus}")
        for i in range(num_gpus):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
    else:
        print("No GPUs available")

def main():
    # 打印可用的 GPU
    print_available_gpus()
    
    # 使用第一个可用的 GPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 示例模型和数据
    model = MyModel().to(device)
    data = torch.randn(10, 3, 224, 224).to(device)  # 示例数据

    # 示例训练循环
    for epoch in range(10):
        # 假设 dataloader 是你的数据加载器
        for batch in dataloader:
            inputs, labels = batch
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            # 其余训练步骤...
    
if __name__ == "__main__":
    main()

3. 结合 CUDA_VISIBLE_DEVICES 和 PyTorch 代码

通过在命令行中设置 CUDA_VISIBLE_DEVICES 环境变量，并在代码中使用 torch.cuda 相关的方法，可以确保你的脚本只使用指定的 GPU。

例如，在命令行中设置环境变量后运行脚本：

bash 复制代码 
CUDA_VISIBLE_DEVICES=0,1 python main.py

然后在 Python 脚本中：

python 复制代码 
import torch

def main():
    # 检查可用的 GPU
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        print(f"Using {num_gpus} GPUs")
        for i in range(num_gpus):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
    else:
        print("No GPUs available, using CPU")
    
    # 使用第一个可用的 GPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # 示例模型和数据
    model = MyModel().to(device)
    data = torch.randn(10, 3, 224, 224).to(device)  # 示例数据

    # 示例训练循环
    for epoch in range(10):
        # 假设 dataloader 是你的数据加载器
        for batch in dataloader:
            inputs, labels = batch
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            # 其余训练步骤...
    
if __name__ == "__main__":
    main()
相关推荐
学步_技术几秒前
Python编码系列—Python抽象工厂模式：构建复杂对象家族的蓝图
开发语言·python·抽象工厂模式
DisonTangor32 分钟前
阿里通义千问开源Qwen2.5系列模型：Qwen2-VL-72B媲美GPT-4
人工智能·计算机视觉
豆浩宇32 分钟前
Halcon OCR检测 免训练版
c++·人工智能·opencv·算法·计算机视觉·ocr
Narutolxy34 分钟前
Python 单元测试：深入理解与实战应用20240919
python·单元测试·log4j
LLSU1336 分钟前
聚星文社AI软件小说推文软件
人工智能
JackieZhengChina39 分钟前
吴泳铭：AI最大的想象力不在手机屏幕，而是改变物理世界
人工智能·智能手机
ShuQiHere40 分钟前
【ShuQiHere】 探索数据挖掘的世界：从概念到应用
人工智能·数据挖掘
嵌入式杂谈41 分钟前
OpenCV计算机视觉：探索图片处理的多种操作
人工智能·opencv·计算机视觉
时光追逐者42 分钟前
分享6个.NET开源的AI和LLM相关项目框架
人工智能·microsoft·ai·c#·.net·.netcore
东隆科技42 分钟前
PicoQuant公司：探索铜铟镓硒（CIGS）太阳能电池技术，引领绿色能源革新
人工智能·能源
热门推荐
01RAG 实践- Ollama+RagFlow 部署本地知识库02组基轨迹建模 GBTM的介绍与实现（Stata 或 R）032024年高教社杯数学建模国赛C题超详细解题思路分析04Coze扣子平台完整体验和实践（附国内和国际版对比）05苍穹外卖面试总结06【2024数模国赛赛题思路公开】国赛B题思路丨附可运行代码丨无偿自提07yolov8实战第五天——yolov8+ffmpg实时视频流检测并进行实时推流——（推流，保姆教学）08CCF-CSP认证考试 202406-3 文本分词 100分题解09【2024高教社杯全国大学生数学建模竞赛】B题 生产过程中的决策问题——解题思路 代码 论文10浏览器插件——ModHeader 的分享和学习