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

torch.cuda.set_device() 函数可以用来设置当前使用的 GPU 设备。如果系统中有多个 GPU 设备，可以通过该函数来指定使用哪一个 GPU。

以下是一个示例，说明如何使用 torch.cuda.set_device() 函数来指定多个 GPU 设备：

import torch

# 指定要使用的 GPU 设备的编号
device_ids = [0, 1]

# 创建一个模型，并将模型移动到指定的 GPU 设备上
model = MyModel().cuda(device_ids[0])
model = torch.nn.DataParallel(model, device_ids=device_ids)

# 创建一个张量，并将张量移动到指定的 GPU 设备上
x = torch.randn(10, 10).cuda(device_ids[0])

# 使用 torch.cuda.set_device() 函数来切换当前使用的 GPU 设备
torch.cuda.set_device(device_ids[1])

# 对张量进行操作
y = model(x)

# 使用 torch.cuda.set_device() 函数切换回原来的 GPU 设备
torch.cuda.set_device(device_ids[0])

在上面的示例中，首先创建了一个包含两个 GPU 设备编号的列表 `device_ids`。接着，创建了一个模型 `MyModel()`，并将该模型移动到列表中的第一个 GPU 设备上。然后，使用 `torch.nn.DataParallel` 将模型复制到其他 GPU 设备上。接着，创建了一个张量 `x`，并将该张量移动到列表中的第一个 GPU 设备上。

在对张量 `x` 进行操作之前，使用 `torch.cuda.set_device()` 函数将当前使用的 GPU 设备切换到列表中的第二个 GPU 设备上。在操作完成后，再将当前使用的 GPU 设备切换回列表中的第一个 GPU 设备上。