DataParallel怎么将输入数据迁移到GPU？

在PyTorch中，DataParallel 是一个模块，它允许你在多个GPU上分布式地训练模型。为了让输入数据自动移动到GPU，你需要做的是将模型和数据包装在 DataParallel 中。以下是一个简单的例子，展示如何使用 DataParallel：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.nn import DataParallel
 
# 假设你有一个模型和一个数据加载器
model = nn.Linear(10, 10)
data_loader = DataLoader(torch.randn(100, 10), batch_size=10)
 
# 在GPU(s)上创建模型的实例
device_ids = [0, 1]  # GPU的ID，例如，0和1表示使用第一个和第二个GPU
model = DataParallel(model, device_ids=device_ids)
 
# 将模型和数据加载器移动到GPU
model.to(torch.device('cuda:0'))  # 选择第一个GPU作为主GPU
 
# 训练循环
for data in data_loader:
    input = data.to(torch.device('cuda:0'))  # 将输入移动到主GPU
    output = model(input)
    loss = loss_function(output, target)
    loss.backward()
    optimizer.step()

在这个例子中，DataParallel 会自动将batch数据均分到所有指定的GPU上，并在前向传播和反向传播时进行all-reduce操作。你只需要将输入数据通过 .to(torch.device('cuda:0')) 移动到主GPU（即 DataParallel 中的第一个GPU），模型和损失函数会自动处理分布式计算。