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

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

相关文章

Python pandas实现excel工作表合并功能详解

Python pandas实现excel工作表合并功能详解

import os,pandas as pd,re #1.获取文件夹下要合并的文件名 dirpath = '文件夹地址' #工作表3特殊处理 需要开始下标和结束下标 begin =...

python中数组和矩阵乘法及使用总结(推荐)

Matrix是Array的一个小的分支,包含于Array。所以matrix 拥有array的所有特性。 但在数组乘和矩阵乘时,两者各有不同,如果a和b是两个matrices,那么a*b,...

TensorFlow在MAC环境下的安装及环境搭建

TensorFlow在MAC环境下的安装及环境搭建

给大家分享一下TensorFlow在MAC系统中的安装步骤以及环境搭建的操作流程。 TensorFlow 底层的图模型结构清晰,容易改造;支持分布式训练;可视化效果好。如果做长期项目,接...

在python中的socket模块使用代理实例

说socket代理之前,先来说说http代理,python的urllib2是自带http代理功能的,可以用如下代码实现:复制代码 代码如下:proxy_handler = urllib2...

python3操作mysql数据库的方法

软硬件环境 OS X EI Capitan Python 3.5.1 mysql 5.6 前言 在开发中经常涉及到数据库的使用,而python对于数据库也有多种解决方法。本文以pyth...