举例说明 torch.cuda.set_device() 如何指定多张GPU

torch.cuda.set_device() 函数用于设置当前使用的 GPU，可以指定一张或多张 GPU。下面是一个指定多张 GPU 的例子：

import torch

# 指定使用的 GPU
device_ids = [0, 1, 2, 3]  # 假设有四张 GPU
torch.cuda.set_device(device_ids[0])  # 设置当前使用的第一张 GPU

# 定义模型并将模型放到多张 GPU 上
model = torch.nn.DataParallel(model, device_ids=device_ids)
model.cuda()

在上面的例子中，我们通过 `device_ids` 列表指定了要使用的多张 GPU，然后将第一张 GPU 设置为当前使用的 GPU。接着，我们将模型使用 `torch.nn.DataParallel` 并指定 `device_ids` 参数来实现在多张 GPU 上并行计算。最后，我们将模型移动到 GPU 上。这样，我们就可以在多张 GPU 上并行计算模型了。

相关问题

torch.cuda.set_device()方法

引用\[1\]和\[2\]介绍了使用torch.cuda.set_device()方法来指定使用的显卡。这个方法可以用来设置主显卡或者指定特定的显卡。在使用这个方法之前，需要先导入torch库。然后可以使用torch.cuda.set_device()方法来设置要使用的显卡的索引。例如，torch.cuda.set_device(1)表示使用索引为1的显卡。接下来，可以使用.cuda()方法将张量或模型加载到指定的显卡上。例如，x = torch.tensor(\[\[1,2,3\],\[4,5,6\]\]).cuda()将张量x加载到当前指定的显卡上。最后，可以使用print(x.device)来打印出张量所在的设备，即显卡的信息。

需要注意的是，引用\[3\]提到了使用torch.cuda.set_device()方法的优先级较低，并且官方文档不建议使用该方法。通常情况下，我们可以直接使用.cuda()方法将模型和数据加载到对应的GPU上，而不需要显式地调用torch.cuda.set_device()方法。
#### 引用[.reference_title]
- *1* *2* *3* [torch.cuda常用指令](https://blog.csdn.net/qq_43307074/article/details/127628498)[target="_blank" data-report-click={"spm":"1018.2226.3001.9630","extra":{"utm_source":"vip_chatgpt_common_search_pc_result","utm_medium":"distribute.pc_search_result.none-task-cask-2~all~insert_cask~default-1-null.142^v91^insertT0,239^v3^insert_chatgpt"}} ] [.reference_item]
[ .reference_list ]

torch.cuda.set_device(device) 改成单机多卡

如果要在单机上使用多个GPU进行计算，可以使用`DataParallel`或者`DistributedDataParallel`这两个PyTorch提供的模块。其中，`DataParallel`将数据分割成多个部分，在多个GPU上分别计算，最后再将结果合并；`DistributedDataParallel`则更加复杂，需要手动进行进程和通信的设置，但是在性能方面更加出色。

使用`DataParallel`的话，可以这样修改：

import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.utils.data.distributed as distributed

device_ids = [0, 1]  # 指定要使用的GPU设备编号
model = nn.DataParallel(model, device_ids=device_ids)  # 使用DataParallel进行多GPU计算

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 训练循环
for epoch in range(num_epochs):
    # 从数据集中加载数据
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    # 训练模型
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

使用`DistributedDataParallel`的话，可以这样修改：

import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
import torch.utils.data.distributed as distributed

# 初始化进程组
dist.init_process_group(backend='nccl', init_method='env://')

# 设置分布式计算相关参数
world_size = dist.get_world_size()
rank = dist.get_rank()
torch.cuda.set_device(rank)

# 加载模型，使用DistributedDataParallel进行多GPU计算
model = MyModel()
model = nn.parallel.DistributedDataParallel(model)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 从数据集中加载数据
train_sampler = distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True, sampler=train_sampler)

# 训练循环
for epoch in range(num_epochs):
    train_sampler.set_epoch(epoch)
    for inputs, labels in train_loader:
        inputs, labels = inputs.cuda(), labels.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

需要注意的是，使用`DistributedDataParallel`时需要先初始化进程组，设置分布式计算相关参数，以及使用`DistributedSampler`对数据进行分布式采样。