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

要在多个GPU上训练PyTorch模型，可以使用`nn.DataParallel`模块。下面是一个使用`nn.DataParallel`模块进行多GPU并行训练的示例：

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

# 定义模型
class TransformerModel(nn.Module):
    def __init__(self):
        super(TransformerModel, self).__init__()
        # ...
    
    def forward(self, x):
        # ...
        return output

# 初始化模型和数据加载器
model = TransformerModel()
train_data = DataLoader(...)
val_data = DataLoader(...)

# 指定设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.device_count() > 1:
    print("Using", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

# 将模型和数据加载器移动到设备上
model.to(device)
train_data.to(device)
val_data.to(device)

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

# 训练模型
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(train_data):
        # 将输入和标签移动到设备上
        inputs = inputs.to(device)
        labels = labels.to(device)
        
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    # 在验证集上测试模型
    with torch.no_grad():
        total_loss = 0
        for inputs, labels in val_data:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            total_loss += loss.item()
        avg_loss = total_loss / len(val_data)
        print("Epoch", epoch, "Validation Loss:", avg_loss)

在上面的代码中，如果检测到有多个GPU可用，就使用`nn.DataParallel`对模型进行并行化。`nn.DataParallel`将自动将数据分配到可用的GPU上，并将梯度聚合在一起。

请注意，如果使用`nn.DataParallel`，则需要将模型和数据加载器移动到设备上，而不是将它们传递给`nn.DataParallel`。