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

你可以使用PyTorch的`torch.nn.DataParallel`模块来实现多GPU并行训练。具体实现步骤如下：

1. 使用`torch.cuda.device_count()`函数检测当前机器上可用的GPU数量，并将模型放到所有可用的GPU上。

2. 使用`torch.nn.DataParallel`将模型封装起来，使其能够在多个GPU上并行运行。

下面是一个简单的示例代码：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# 定义模型和数据加载器
model = MyModel().cuda()
data_loader = DataLoader(dataset)

# 检测可用的GPU数量，并将模型放到所有可用的GPU上
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

# 定义优化器和损失函数
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 开始训练
for epoch in range(num_epochs):
    for batch in data_loader:
        inputs, labels = batch
        inputs = inputs.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

在这个示例代码中，如果机器上有多个GPU，`nn.DataParallel`会自动使用所有可用的GPU来并行运行模型。注意，在使用`nn.DataParallel`时，模型的输入和输出需要是`torch.Tensor`类型，不能是Python原生数据类型，否则会报错。