pytorch指定用多张显卡训练_Pytorch中多GPU训练指北

165 阅读 0 评论 109 点赞

我是靠谱客的博主勤恳指甲油，这篇文章主要介绍pytorch指定用多张显卡训练_Pytorch中多GPU训练指北，现在分享给大家，希望可以做个参考。

前言

在数据越来越多的时代，随着模型规模参数的增多，以及数据量的不断提升，使用多GPU去训练是不可避免的事情。Pytorch在0.4.0及以后的版本中已经提供了多GPU训练的方式，本文简单讲解下使用Pytorch多GPU训练的方式以及一些注意的地方。

这里我们谈论的是单主机多GPUs训练，与分布式训练不同，我们采用的主要Pytorch功能函数为DataParallel而不是DistributedParallel，后者为多主机多GPUs的训练方式，但是在实际任务中，两种使用方式也存在一部分交集。

使用方式

使用多卡训练的方式有很多，当然前提是我们的设备中存在两个及以上的GPU：使用命令nvidia-smi查看当前Ubuntu平台的GPU数量(Windows平台类似)，其中每个GPU被编上了序号：[0,1]：

在我们设备中确实存在多卡的条件下，最简单的方法是直接使用torch.nn.DataParallel将你的模型wrap一下即可：

net=torch.nn.DataParallel(model)

这时，默认所有存在的显卡都会被使用。

如果我们机子中有很多显卡(例如我们有八张显卡)，但我们只想使用0、1、2号显卡，那么我们可以：

net=torch.nn.DataParallel(model,device_ids=[0,1,2])

或者这样：

本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。