如何使用GPU以及Pytorch中GPU的相关操作

最新推荐文章于 2024-09-13 08:59:18 发布

原创最新推荐文章于 2024-09-13 08:59:18 发布 · 555 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#pytorch #人工智能 #python

本文介绍了如何在PyTorch中检查GPU状态、设置使用特定GPU、数据并行和单卡运行的方法，以及处理多显卡训练时的常见问题，包括nn.DataParallel的使用和CUDA_VISIBLE_DEVICES环境变量的设置。

一、GPU的基本信息

1.查看cuda是否可用：torch.cuda.is_available()

2.查看gpu数量：torch.cuda.device_count()

3.查看gpu名字，设备索引默认从0开始：torch.cuda.get_device_name(0)

4.当前设备的索引：torch.cuda.current_device()

5.查看gpu的使用情况：nvidia-smi

6.每隔1s刷新一次gpu的使用情况：watch -n 1 nvidia-smi

二、代码中，如何设定使用哪种gpu

1.单卡的时候，没有选择的余地，就一张

2.多卡的时候，分两种情况，一个是数据并行，多张卡一起工作；另一个是只在一张卡上运行，比如由4张卡[0,1,2,3]，我想在卡1上运行任务。

情况1：数据并行

#配置device_ids，选择你想用的卡编号。
device_ids= [0, 1, 2]
if torch.cuda.device_count() >1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model=torch.nn.DataParallel(model, device_ids)

此时的gpyu使用情况：

+-----------------------------------------------------------------------------+
| Processes: GPU Memory |
| GPU PID Type Process name Usage |
|=============================================================================|
| 0 19272 C python 9009MiB |
| 1 19272 C python 5753MiB |
| 2 19272 C python 5753MiB |
| 3 19272 C python 5755MiB |
+-----------------------------------------------------------------------------+
模型的结构默认就是在device_ids[0]，即第一块卡上，也就解释了为什么第一块卡的显存会占用的比其他卡要更多一些。进一步说也就是当你调用nn.DataParallel的时候，只是在你的input数据是并行的，但是你的output loss却不是这样的，每次都会在第一块GPU相加计算，这就造成了第一块GPU的负载远远大于剩余其他的显卡。

情况二：一张卡上运行
类似tensorflow指定GPU的方式，使用CUDA_VISIBLE_DEVICES

1.1 直接终端中设定：

CUDA_VISIBLE_DEVICES=1 python main.py

1.2 python代码中设定：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

2.使用函数 set_deviceimport torch
torch.cuda.set_device(id)
分布式运行训练程序的问题
8881服务器上四张卡已经有一张再跑，想要只使用1，2，3号卡

此处为多显卡训练设置

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3"
device_ids = [1,2,3]
model = nn.DataParallel(model, device_ids) # multi-gpu training
运行时报错：

RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found
问题就出在device处，因为下面这条语句，默认将模型放在了cuda 0上，也就是我们没有使用的显卡0。于是会产生错误

device = torch.device("cuda")
修改后成功

device = torch.device("cuda:1"）