在使用 PyTorch 框架时,可以通过以下步骤查看可用的 GPU 数量,指定使用的 GPU 编号,并在代码中体现这一点。下面以2个GPU为例:
目录
- 一、脚本代码块实现
- 1. 查看可用的 GPU
- 2. 指定使用 GPU 的数量和编号
- 使用 `CUDA_VISIBLE_DEVICES` 环境变量
- 将模型和数据移动到指定的 GPU
- 使用 `DataParallel` 进行多 GPU 并行计算
- 3. 示例完整代码
- 二、命令行实现
- 1. 使用 `CUDA_VISIBLE_DEVICES` 环境变量
- 只使用 GPU 0
- 使用 GPU 0 和 GPU 1
- 2. 在 Python 代码中自动检测可用 GPU
- 3. 结合 `CUDA_VISIBLE_DEVICES` 和 PyTorch 代码
一、脚本代码块实现
1. 查看可用的 GPU
import torchif 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:
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。例如:
import torch# 检查是否有 GPU 可用
if torch.cuda.is_available():# 将模型移动到 GPU 0device0 = torch.device('cuda:0')model = MyModel().to(device0)# 将数据移动到 GPU 1device1 = 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 上并行运行:
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,并在代码中体现:
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
CUDA_VISIBLE_DEVICES=0 python main.py
使用 GPU 0 和 GPU 1
CUDA_VISIBLE_DEVICES=0,1 python main.py
2. 在 Python 代码中自动检测可用 GPU
在 Python 代码中,可以自动检测可用的 GPU,并根据需要配置模型和数据。以下是示例代码:
import torchdef 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():# 打印可用的 GPUprint_available_gpus()# 使用第一个可用的 GPUdevice = 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 = batchinputs, 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。
例如,在命令行中设置环境变量后运行脚本:
CUDA_VISIBLE_DEVICES=0,1 python main.py
然后在 Python 脚本中:
import torchdef main():# 检查可用的 GPUif 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")# 使用第一个可用的 GPUdevice = 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 = batchinputs, labels = inputs.to(device), labels.to(device)outputs = model(inputs)# 其余训练步骤...if __name__ == "__main__":main()