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

目录

一、脚本代码块实现1. 查看可用的 GPU2. 指定使用 GPU 的数量和编号使用 `CUDA_VISIBLE_DEVICES` 环境变量将模型和数据移动到指定的 GPU使用 `DataParallel` 进行多 GPU 并行计算

3. 示例完整代码

二、命令行实现1. 使用 `CUDA_VISIBLE_DEVICES` 环境变量只使用 GPU 0使用 GPU 0 和 GPU 1

2. 在 Python 代码中自动检测可用 GPU3. 结合 `CUDA_VISIBLE_DEVICES` 和 PyTorch 代码

一、脚本代码块实现

1. 查看可用的 GPU

<code>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：

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 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 上并行运行：

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 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。

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

CUDA_VISIBLE_DEVICES=0,1 python main.py

然后在 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()