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

相关文章

Python中base64与xml取值结合问题

Base64是一种用64个字符来表示任意二进制数据的方法。 用记事本打开exe、jpg、pdf这些文件时,我们都会看到一大堆乱码,因为二进制文件包含很多无法显示和打印的字符,所以,如果要...

简单了解Python下用于监视文件系统的pyinotify包

什么是inotify:   Inotify是一个事件驱动的通知机制,Inotify 提供一个简单的API,使用最小的文件描述符,并且允许细粒度监控。与 inotify 的...

在Python中执行系统命令的方法示例详解

前言 Python经常被称作“胶水语言”,因为它能够轻易地操作其他程序,轻易地包装使用其他语言编写的库。在Python/wxPython环境下,执行外部命令或者说在Python程序中启动...

Python中str.format()详解

1. str.format 的引入 在 Python 中,我们可以使用 + 来连接字符串,在简单情况下这种方式能够很好的工作。但是当我们需要进行复杂的字符串连接时,如果依然使用 + 来...

Python实现简易端口扫描器代码实例

Python实现简易端口扫描器代码实例

在网上的一些资料的基础上自己又添了些新内容,算是Python socket编程练手吧。 #coding=utf-8 import socket import time import...