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)
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐