如下所示: device = torch.device(cuda:0 if torch.cuda.is_available() else cpu)#第一行代码model.to(device)#第二行代码 首先是上面两行代码放在读取数据之前。 mytensor = my_tensor.to(device)#第三行代码 然后是第三行代码。这句代码的意思是将所有最
如下所示:
首先是上面两行代码放在读取数据之前。
然后是第三行代码。这句代码的意思是将所有最开始读取数据时的tersor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。 需要注意的是这句话并不像前面的两行代码一样只需要写一遍,第三行代码需要写的次数就等于需要保存到GPU上的tensor变量个数;一般情况下这些tensor变量都是最开始读取数据时的tensor变量,后面所衍生的变量自然也都在GPU之上。 以上是使用单个GPU的情况。当你拥有多个GPU时,要想使用多个GPU进行训练和测试,需要在第一二行代码之间插上下面这样一个判断语句,其余的写法也都是一样的。
使用多个GPU的原理就是通过上面这句代码将model在每个GPU上分别保存一份,然后对model的输入tensor进行自动的分割,每个GPU计算tensor的一部分,这样就能实现计算量的平均分配。 在每个model计算完成之后,DataParallel将这些结果进行收集和融合,之后再将结果返回。 |
2019-06-18
2019-07-04
2021-05-23
2021-05-27
2021-05-27