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实现基本数据结构中栈的操作示例

Python实现基本数据结构中栈的操作示例

本文实例讲述了Python实现基本数据结构中栈的操作。分享给大家供大家参考,具体如下: #! /usr/bin/env python #coding=utf-8 #Python实现基...

python实现文本界面网络聊天室

python实现文本界面网络聊天室

Hello大家好,今天说一下python的socket编程,基于python的socket通信的文本框网络聊天 首先,实验环境: 一个云服务器(我们这里是用的阿里云,大家将就自己的条件吧...

python二分法实现实例

1.算法:(设查找的数组期间为array[low, high]) (1)确定该期间的中间位置K(2)将查找的值T与array[k]比较。若相等,查找成功返回此位置;否则确定新的查找区域,...

深入理解Python中的super()方法

前言 python的类分别有新式类和经典类,都支持多继承。在类的继承中,如果你想要重写父类的方法而不是覆盖的父类方法,这个时候我们可以使用super()方法来实现 python语言与C+...

Django学习笔记之ORM基础教程

Django学习笔记之ORM基础教程

ORM简介 ORM概念 对象关系映射(Object Relational Mapping,简称ORM)模式是一种为了解决面向对象与关系数据库存在的互不匹配的现象的技术。 简单的说,OR...