深度学习方法——pytorch下GPU与CPU调用的切换_cuda转cpu

一、准备工作

在设置GPU调用之前，我们首先需要确认当前设备的显卡是否支持cuda编程，否则所有基于cuda的操作都会报错。我们可以通过torch.cuda.is_available()来进行确认，并通过torch.cuda.get_device_name()查看当前可用的显卡设备名称：

import torch
print(torch.cuda.get_device_name(0)) # 查看使用的设备名称
torch.cuda.is_available() 		     # 验证cuda是否正常安装并能够使用
1
2
3

输出：
NVIDIA GeForce RTX 3070 Laptop GPU
True

---

当我们有多块显卡时，可以通过修改系统环境变量中的"CUDA_VISIBLE_DEVICES"，用于指定可以使用的显卡：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'  		# 只允许使用0号显卡
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2,3,4'  # 允许使用多块显卡
1
2
3

PS：当前博主的电脑只含一块显卡所以无法验证这个代码声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】