pytorch指定用多张显卡训练_pytorch多gpu并行训练

可以使用PyTorch中的DataParallel模块来指定使用多张显卡进行并行训练。以下是一个简单的示例代码：

import torch.nn as nn
import torch.optim as optim
import torch.utils.data
from torch.utils.data import DataLoader
from torch.nn.parallel import DataParallel

# 定义模型
model = nn.Sequential(
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 1)
)

# 定义优化器和损失函数
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# 将模型放到多个GPU上
model = nn.DataParallel(model)

# 加载数据
train_loader = DataLoader(dataset)

# 训练模型
for epoch in range(num_epochs):
    for data in train_loader:
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        # 正向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

在这个示例中，我们使用了DataParallel模块将模型放到多个GPU上，然后使用DataLoader加载数据进行训练。在正向传播和反向传播过程中，我们需要将输入数据和标签数据都移动到GPU上进行计算。