关于 torch.nn.DataParallel 的问题

我是深度学习领域的新手。现在我正在复制一篇论文的代码。因为他们使用多个 GPU，所以有一个命令torch.nn.DataParallel(model, device_ids= args.gpus).cuda()在代码中。但我只有一个GPU，什么 我应该更改此代码以匹配我的 GPU 吗？

谢谢你！

DataParallel也应该在单个 GPU 上工作，但您应该检查是否args.gpus仅包含要使用的设备的 ID（应为 0）或None。 选择None将使模块使用所有可用的设备。

你也可以删除DataParallel因为您不需要它，只需通过调用即可将模型移动到 GPUmodel.cuda()或者，如我所愿，model.to(device) where device是设备的名称。

Example:

此示例展示了如何在单个 GPU 上使用模型，并使用以下命令设置设备.to()代替.cuda().

from torch import nn
import torch

# Set device to cuda if cuda is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create model
model = nn.Sequential(
  nn.Conv2d(1,20,5),
  nn.ReLU(),
  nn.Conv2d(20,64,5),
  nn.ReLU()
)

# moving model to GPU
model.to(device)

如果你想使用DataParallel你可以这样做

# Optional DataParallel, not needed for single GPU usage
model1 = torch.nn.DataParallel(model, device_ids=[0]).to(device)
# Or, using default 'device_ids=None'
model1 = torch.nn.DataParallel(model).to(device)