官方教程 首先将神经网络转到GPU上,前提条件是CUDA可以用。所以我们需要进行检测。
import torch device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device)如果CUDA可用输出cuda:0,不可用则输出cpu
这里我们使用randn创建随机数据集,并将这些数据使用DataLoader封装成30大小的Tensor用于训练。 该接口的目的:将自定义的Dataset根据batch size大小、是否shuffle等封装成一个Batch Size大小的Tensor,用于后面的训练。
#创建一个随机数据集,制作假数据 class RandomDataset(Dataset): def __init__(self, size, length): self.len = length self.data = torch.randn(length, size) def __getitem__(self, index): return self.data[index] def __len__(self): return self.len #随机数集,shffule=True在表示不同批次的数据遍历时,打乱顺序,每份30个数据 rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size), batch_size=batch_size, shuffle=True)这里我们定义一个很简单的神经网络,模型仅获得输入,执行线性运算并给出输出。在模型中放置了一条打印语句,以监视输入和输出张量的大小。请注意批次等级0上打印的内容。
class Model(nn.Module): #定义net的初始化函数,这个函数定义了该神经网络的基本结构 def __init__(self, input_size, output_size): #复制并使用Net的父类的初始化方法,即先运行nn.Module的初始化函数 super(Model, self).__init__() #定义fc(fullconnect)全连接函数为线性函数:y = Wx + b,并将5个节点连接到2个节点上。 self.fc = nn.Linear(input_size, output_size) #定义该神经网络的向前传播函数,该函数必须定义,一旦定义成功,向后传播函数也会自动生成(autograd) def forward(self, input): #输入input经过全连接函数fc后,更新output output = self.fc(input) #打印input和output的大小 print("\t In Model: input size", input.size(), "output size", output.size()) return output新建网络并实现数据并行化,PyTorch默认只使用一个GPU,可以使用DataParallel来实现使用多个GPUmodel = nn.DataParallel(model) 调用model.to(device)返回的是一个model新的备份,而不是重写了model. 需要将其分配到一个新的张量,并在GPU上使用这个张量.
#建一个模型实例 model = Model(input_size, output_size) #是否有多个GPU if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs #PyTorch默认只使用一个GPU,使用DataParallel来实现使用多个GPU model = nn.DataParallel(model) else:print("only one") #将模型放入GPU model.to(device)这里遇到了问题,我的电脑明明有两个GPU,但是torch.cuda.device_count()只显示有一个。 运行后输出
only one但是任务管理器中显示我有两个GPU 待解决。相同问题
具体代码和注释如下
#遍历我们随机产生的数据集 for data in rand_loader: #将数据放入GPU中,赋值给input input = data.to(device) #将input投入模型更新output output = model(input) #打印input和output的大小 print("Outside: input size", input.size(), "output_size", output.size())在GPU上训练神经网络时,需要在网络模型定义好后放入GPU中net.to(device),这个方法将递归遍历所有模块,并将其参数和缓冲区都转换为CUDA张量。然后要将每一步的输入和数据都放入GPU中,比如在训练网络时,每一次获取数据后,都要inputs, labels = data[0].to(device), data[1].to(device)。
本文模型运行结果(只有一个GPU)
D:\Python\python.exe C:/Users/465/Desktop/try/通信/test.py only one In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2]) Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])30是batch_size = 30每份是30个数据,5和2是self.fc = nn.Linear(input_size, output_size)将5个节点连接到了2个节点上。总共有data_size = 100个数据。
如果有两个GPU
# on 2 GPUs Let's use 2 GPUs! In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2]) Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2]) In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2]) In Model: input size torch.Size([5, 5]) output size torch.Size([5, 2]) Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])DataParallel会自动分割数据,并将作业订单发送到多个GPU上的多个模型。每个模型完成工作后,DataParallel会收集并合并结果,然后再将结果返回。