举例说明 torch.cuda.set_device() 如何指定多张GPU
时间: 2024-02-03 11:15:56 浏览: 490
pytorch 指定gpu训练与多gpu并行训练示例
torch.cuda.set_device() 函数用于设置当前使用的 GPU,可以指定一张或多张 GPU。下面是一个指定多张 GPU 的例子:
```python
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 上并行计算模型了。
阅读全文