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设计】。

相关文章

Win10 安装PyCharm2019.1.1(图文教程)

Win10 安装PyCharm2019.1.1(图文教程)

使用Python进行开发时,已经安装好Python,我们可以在终端使用Python,也可以用文本编辑器编写代码后然后用Python执行代码,但这些都不是很方便,也不是很智能和不够友好,需...

python3 发送任意文件邮件的实例

实例如下所示: #!/usr/bin/python # -*- coding: UTF-8 -*- import smtplib import email.mime.multipar...

浅谈python 线程池threadpool之实现

首先介绍一下自己使用到的名词: 工作线程(worker):创建线程池时,按照指定的线程数量,创建工作线程,等待从任务队列中get任务; 任务(requests):即工作线程处理的任务,任...

Python While循环语句实例演示及原理解析

Python While循环语句实例演示及原理解析

这篇文章主要介绍了Python While循环语句实例演示及原理解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下 Python 编程...

对Python中创建进程的两种方式以及进程池详解

在Python中创建进程有两种方式,第一种是: from multiprocessing import Process import time def test(): whil...