利用DataParallel()进行Python中数据并行计算的示例代码

DataParallel()是PyTorch库中的一个函数，用于实现数据并行计算。它能够自动将模型和数据划分到多个GPU上进行并行计算，从而提高训练或推理的速度。下面是利用DataParallel()进行数据并行计算的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DataParallel

# 定义网络模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = MyModel()
# 将模型包装到DataParallel中
model = DataParallel(model)

# 设置输入数据
inputs = torch.randn(1000, 10)
# 设置目标值
targets = torch.randn(1000, 10)

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

# 训练模型
for epoch in range(100):
    # 将输入数据和目标值发送到GPU上
    inputs = inputs.cuda()
    targets = targets.cuda()

    # 前向传播
    outputs = model(inputs)

    # 计算损失
    loss = criterion(outputs, targets)

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

print("训练完成！")

在上面的示例中，首先定义了一个简单的线性模型MyModel，接着创建了模型实例，并将其包装到DataParallel中以实现数据并行计算。然后，创建了输入数据和目标值，定义了损失函数和优化器。接下来，进行训练过程，在每个epoch中将输入数据和目标值发送到GPU上，调用模型进行前向传播，并计算损失，然后完成反向传播和优化。最后，打印出训练完成的消息。

需要注意的是，在使用DataParallel进行数据并行计算时，会根据多个GPU自动将数据划分到不同的GPU上，并在每个GPU上分别计算。因此，在使用DataParallel时，不需要手动编写代码进行数据的划分和传输，这部分工作由DataParallel自动完成。

除了训练过程，DataParallel同样适用于推理过程。只需将输入数据发送到GPU上，然后调用模型进行前向传播即可。

希望以上示例代码能够帮助你理解如何使用DataParallel()进行数据并行计算。