关于火炬.nn.DataP

Question

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

谢谢你！在

Answer 1

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

您也可以删除DataParallel，因为您不需要它，只需调用model.cuda()或者，我更喜欢的是，model.to(device)，其中device是设备的名称。在

示例：

这个例子展示了如何在单个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，可以这样做

^{pr2}$