pytorch 使用单个GPU与多个GPU进行训练与测试的方法

yipeiwu_com5年前Python基础

如下所示:

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进行训练和测试,需要在第一二行代码之间插上下面这样一个判断语句,其余的写法也都是一样的。

if torch.cuda.device_count() > 1:
 model = nn.DataParallel(model)

使用多个GPU的原理就是通过上面这句代码将model在每个GPU上分别保存一份,然后对model的输入tensor进行自动的分割,每个GPU计算tensor的一部分,这样就能实现计算量的平均分配。在每个model计算完成之后,DataParallel将这些结果进行收集和融合,之后再将结果返回。

以上这篇pytorch 使用单个GPU与多个GPU进行训练与测试的方法就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

Face++ API实现手势识别系统设计

Face++ API实现手势识别系统设计

        通过普通摄像头拍摄出的照片来进行识别是存在很大的困难的,但是有困难才能找到更好的方法去解决。在百度上...

Python实现删除列表中满足一定条件的元素示例

本文实例讲述了Python实现删除列表中满足一定条件的元素。分享给大家供大家参考,具体如下: 从列表中删除满足一定条件的元素。 如:删除一个列表中长度为0的元素,或者删除列表中同时是2和...

rabbitmq(中间消息代理)在python中的使用详解

rabbitmq(中间消息代理)在python中的使用详解

在之前的有关线程,进程的博客中,我们介绍了它们各自在同一个程序中的通信方法。但是不同程序,甚至不同编程语言所写的应用软件之间的通信,以前所介绍的线程、进程队列便不再适用了;此种情况便只能...

Python 实现毫秒级淘宝抢购脚本的示例代码

本篇文章主要介绍了Python 通过selenium实现毫秒级自动抢购的示例代码,通过扫码登录即可自动完成一系列操作,抢购时间精确至毫秒,可抢加购物车等待时间结算的,也可以抢聚划算的商品...

python画出三角形外接圆和内切圆的方法

python画出三角形外接圆和内切圆的方法

刚看了《最强大脑》中英对决,其中难度最大的项目需要选手先脑补泰森多边形,再找出完全相同的两个泰森多边形。在惊呆且感叹自身头脑愚笨的同时,不免手痒想要借助电脑弄个图出来看看,闲来无事吹吹牛...