pytorch多卡gpu
时间: 2023-11-08 08:05:12 浏览: 102
PyTorch支持使用多卡GPU进行模型训练,可以使用`torch.nn.DataParallel`或`torch.nn.parallel.DistributedDataParallel`来实现。
`torch.nn.DataParallel`会在多个GPU上复制模型,并将输入数据分配到不同的GPU上进行计算,最后将结果合并。而`torch.nn.parallel.DistributedDataParallel`则是使用分布式的方式进行模型训练,每个GPU都有一个独立的进程,通过消息传递的方式进行通信和同步。
使用`torch.nn.DataParallel`的示例代码如下:
```python
import torch
import torch.nn as nn
# 定义模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.fc1 = nn.Linear(128 * 8 * 8, 10)
def forward(self, x):
x = nn.functional.relu(self.conv1(x))
x = nn.functional.max_pool2d(x, 2)
x = nn.functional.relu(self.conv2(x))
x = nn.functional.max_pool2d(x, 2)
x = x.view(-1, 128 * 8 * 8)
x = self.fc1(x)
return x
# 初始化模型和数据
model = MyModel()
data = torch.randn(32, 3, 32, 32)
# 使用DataParallel进行多卡训练
model = nn.DataParallel(model)
output = model(data)
```
阅读全文