pytorch使用多张显卡

在PyTorch中使用多张显卡可以通过数据并行的方式实现模型的训练和推理加速。下面是使用多张显卡的一般步骤：

1. 导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader

2. 定义模型：

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 定义你的模型结构

    def forward(self, x):
        # 定义前向传播过程
        return x

3. 初始化多个GPU设备：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device_ids = [0, 1, 2]  # 假设有3个可用的GPU设备

4. 将模型放到多个GPU上：

model = MyModel()
model = nn.DataParallel(model, device_ids=device_ids)
model.to(device)

5. 定义损失函数和优化器：

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

6. 加载数据并创建数据加载器：

train_dataset = ...
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

7. 训练模型：

model.train()
for epoch in range(num_epochs):
    for data, labels in train_loader:
        data = data.to(device)
        labels = labels.to(device)

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

8. 在推理阶段，将输入数据也放到GPU上：

model.eval()
with torch.no_grad():
    inputs = inputs.to(device)
    outputs = model(inputs)

这样就可以在多张显卡上并行地训练和推理模型了。注意，当使用`nn.DataParallel`时，模型的`forward`函数会自动在多个GPU上运行，而不需要我们手动指定。