nn.DataParallel

nn.DataParallel 是 PyTorch 中的一个模块，用于在多个 GPU 上并行运行模型。当有多个 GPU 并且想要利用它们来加速训练或推理时，这个模块会非常有用。nn.DataParallel 通过对模型中的每个子模块进行复制，并将输入数据分割成多个部分，然后在每个 GPU 上并行处理这些部分来实现并行化。

使用 nn.DataParallel 的基本步骤是：

1. 初始化模型：首先，你需要定义你的神经网络模型。
2. 包装模型：然后，使用 nn.DataParallel来包装你的模型。你需要传递你的模型和一个设备列表作为参数。设备列表通常是一个包含你想要使用的 GPU 的标识符的列表（例如 [0, 1]表示使用第一个和第二个 GPU）。
3. 将数据移动到 GPU：在将数据传递给模型之前，你需要确保数据已经被移动到了 GPU上。这可以通过使用 .to(device) 方法来实现，其中 device 是一个表示目标 GPU 的对象。
4. 前向传播：现在你可以像平常一样调用模型进行前向传播，nn.DataParallel 会在后台自动处理并行化。

import torch
import torch.nn as nn
import torch.nn.parallel

class SimpleModel(nn.Module):
	def __init__(self):
		super(SimpleModel,self).__init__()
		self.fc = nn.Linear(10,10)
	def forward(self,x):
		return self.fc(x)
model = SimpleModel()
if torch.cuda.device_count()>1:
	print(torch.cuda.device_count())
	model = nn.DataParallel(model,device_ids=[0,1])
#将模型移动到GPU上
model.to("cuda")
input_data = torch.randn(16,10).to("cuda")
output_data = model(input_data)