Python 关于torch.nn.DataParallel的几个问题

我在深度学习领域是新手。现在我正在复制一篇论文的代码。由于它们使用多个GPU，因此在代码中有一个命令

torch.nn.DataParallel（model，device_id=args.gpu）.cuda（）

。但我只有一个GPU，什么 我应该更改此代码以匹配我的GPU吗

---

谢谢大家!

DataParallel

也应在单个GPU上工作，但您应检查

args.GPU

是否仅包含要使用的设备的id（应为0）或

None

。 选择

None

将使模块使用所有可用设备

此外，您还可以删除不需要的

DataParallel

，只需调用

model.cuda（）

或我喜欢的

model.to（device）

即可将模型移动到GPU，其中

device

是设备的名称

示例：

此示例演示如何在单个GPU上使用模型，使用

.to（）

而不是

.cuda（）

设置设备

来自火炬导入nn的

进口火炬
#如果cuda可用，则将设备设置为cuda
device=torch.device（“cuda”如果torch.cuda.is_可用（）否则为“cpu”）
#创建模型
模型=nn.顺序(
nn.Conv2d（1,20,5），
nn.ReLU（），
nn.Conv2d（20,64,5），
nn.ReLU（）
)
#将模型移动到GPU
型号.至（设备）

如果您想使用

DataParallel

，您可以这样做

#可选DataParallel，不需要用于单个GPU
model1=torch.nn.DataParallel（model，device_id=[0]）到（device）
#或者，使用默认的“设备\u id=None”
model1=torch.nn.DataParallel（model）.to（设备）