PyTorch——GPU管理

本文主要介绍PyTorch相关的GPU管理

GPU管理整体说明

  • PyTorch的GPU管理包含以下接口：
    • cuda设备可用性判断：torch.cuda.is_available()
    • 设备对象创建和获取：torch.device('cuda')
    • 张量切换设备：tensor.to(gpu_device)
    • 模型切换设备：model.to(gpu_device)
    • 模型参数设备打印：print(f"参数所在设备: {param.device}") for param in model.parameters(): # model.parameters()
    • cuda设备可用数量查看：torch.cuda.device_count()
    • 当前环境cuda设备管理：torch.cuda.device(1)
    • 当前环境cuda设备号查看：torch.cuda.current_device()

Torch的默认指定GPU设定

  • PyTorch 会默认选择一个 GPU 作为当前设备（通常是编号为 0 的 GPU），原因如下：
    • 单 GPU 环境 ：如果系统中只有一块 GPU，PyTorch 会默认使用它（编号为 0）
    • 多 GPU 环境 ：如果有多个 GPU，PyTorch 仍然会默认使用编号为 0 的 GPU，除非显式地指定使用其他 GPU
    • 简化开发 ：默认设备机制使得开发者在不显式指定设备的情况下，代码仍然可以正常运行
  • 默认 GPU 的设定是为了方便开发者，避免每次都需要手动指定设备

GPU管理先关概念

设备编号（Device Index）

  • PyTorch 使用从 0 开始的整数编号来标识 GPU
  • 例如，如果有 4 块 GPU，它们的编号分别是 0、1、2、3
  • 可以通过 torch.cuda.device_count() 获取当前可用的 GPU 数量

当前设备（Current Device）

  • PyTorch 会跟踪当前正在使用的 GPU，称为「当前设备」
  • 默认情况下，当前设备是编号为 0 的 GPU
  • 可以通过 torch.cuda.current_device() 获取当前设备的编号

切换设备简单示例

  • 可以使用 torch.cuda.set_device(device_id) 来切换当前设备

    1
    2
    torch.cuda.set_device(1)  # 切换到编号为 1 的 GPU
    # 接下来的操作会默认在指定GPU上

  • 也可以通过 torch.device 来指定设备：

    1
    2
    device = torch.device("cuda:1")  # 使用编号为 1 的 GPU
    tensor = torch.tensor([1, 2, 3]).to(device)

GPU高阶管理

  • 如果没有显式指定设备，PyTorch 会将张量和模型放在默认 GPU（通常是 cuda:0）上

  • 例如：

    1
    x = torch.tensor([1, 2, 3]).cuda()  # 默认放在 cuda:0 上

  • 可以通过环境变量 CUDA_VISIBLE_DEVICES 控制哪些 GPU 对 PyTorch 可见。例如：

    1
    export CUDA_VISIBLE_DEVICES=1,2  # 只让 GPU 1 和 2 对 PyTorch 可见

    在这种情况下，PyTorch 会将可见的 GPU 重新编号为 0 和 1

代码示例

  • 使用示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    ## 设备管理
    # 创建一个 CPU 设备对象
    cpu_device = torch.device('cpu')
    # 创建一个 GPU 设备对象，如果有多个 GPU，可以指定编号，如 'cuda:1'
    gpu_device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    ## 检查GPU是否可用
    # if torch.cuda.is_available()

    ## 获取GPU的数量
    num_gpus = torch.cuda.device_count()

    ## 张量切换设备
    # 将张量移动到 GPU 设备
    tensor = tensor.to(gpu_device)
    # 或者使用 .cuda() 方法
    tensor = tensor.cuda() if torch.cuda.is_available() else tensor
    # 将张量移动到 CPU 设备
    tensor = tensor.cpu()
    ## 检查张量所在设备
    print(f"张量所在设备: {tensor.device}")

    ## 模型切换设备(跟张量操作一致)
    model = torch.nn.Linear(10, 10)
    # 将模型移动到 GPU 设备
    model = model.to(gpu_device)
    # 或者使用 .cuda() 方法
    model = model.cuda() if torch.cuda.is_available() else model
    # 将模型移动到 CPU 设备
    model = model.cpu()
    ## 检查参数所在设备
    for param in model.parameters(): # model.parameters()
    	print(f"参数所在设备: {param.device}")

    ## 当前环境的多设备管理
    # 指定GPU设备方案一
    if torch.cuda.is_available() and num_gpus > 1:
    	with torch.cuda.device(1):  # 指定使用第二个 GPU
    		tensor = torch.randn(3, 3).cuda()
    		print(f"张量所在设备: {tensor.device}")
    # 指定GPU设备方案二
    if torch.cuda.is_available() and num_gpus > 1:
       	torch.cuda.set_device(1)  # 指定使用第二个 GPU
    # 输出设备号
    if torch.cuda.is_available():
    	current_device = torch.cuda.current_device()
    	print(f"当前使用的 GPU 编号: {current_device}")

    ## 更清晰的设备管理方式
    # 将 tensor 移动到 cuda:0
    gpu_device0 = torch.device('cuda:0')
    tensor0 = tensor.to(gpu_device0)

    # 将 tensor 移动到 cuda:1
    gpu_device1 = torch.device('cuda:1')
    tensor1 = tensor.to(gpu_device1)