如何我想使用多个GPU在device上,怎么办
时间: 2023-08-16 12:06:31 浏览: 49
如果您想在PyTorch中使用多个GPU,可以使用`torch.nn.DataParallel`模块。这个模块可以自动将数据划分成多个块,分布到各个GPU上进行计算。以下是一个示例代码:
```python
import torch
import torch.nn as nn
from torch.nn.parallel import DataParallel
# 定义模型
model = nn.Sequential(
nn.Linear(in_features=10, out_features=5),
nn.ReLU(),
nn.Linear(in_features=5, out_features=1)
)
# 将模型放在多个GPU上
model_parallel = DataParallel(model, device_ids=[0, 1])
# 定义输入数据
input_data = torch.randn(64, 10)
# 前向传播
output = model_parallel(input_data)
# 计算损失
loss = nn.MSELoss()(output, torch.ones_like(output))
# 反向传播
loss.backward()
# 更新参数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
optimizer.step()
```
在上面的示例中,我们首先定义了一个模型,并使用`DataParallel`模块将模型放在两个GPU上。然后,我们定义了输入数据,通过`model_parallel`进行前向传播,计算损失,并使用反向传播对参数进行更新。使用`DataParallel`模块,我们可以轻松地将模型放在多个GPU上进行训练。