Pytorch基础:torch.cuda.set_device函数

相关阅读

Pytorch基础https://blog.csdn.net/weixin_45791458/category_12457644.html?spm=1001.2014.3001.5482


torch.cuda.set_device函数用于设置当前使用的cuda设备,在当拥有多个可用的GPU且能被pytorch识别的cuda设备情况下(环境变量CUDA_VISIBLE_DEVICES可以影响GPU设备到cuda设备的映射)。

由于有些情况下,可以不显式写出cuda设备的编号,此时指的是当前使用的cuda设备,默认为cuda0设备,如下所示。

复制代码
import torch
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备

a=torch.tensor([1, 2, 3, 4, 5], device='cuda') # 在当前cuda设备上创建一个张量
print(a)

device=torch.device('cuda')                    # 创建一个当前cuda设备

b=torch.tensor([1, 2, 3, 4, 5], device=device) # 在当前cuda设备上创建一个张量
print(b)

c=torch.tensor([1, 2, 3, 4, 5]).to(device)     # 将张量移动到当前cuda设备
print(c)

d=torch.tensor([1, 2, 3, 4, 5]).to('cuda')     # 将张量移动到当前cuda设备
print(d)

e=torch.tensor([1, 2, 3, 4, 5]).cuda()         # 将张量移动到当前cuda设备
print(e)

f=torch.tensor([1, 2, 3, 4, 5]).cuda(device)   # 将张量移动到当前cuda设备
print(f)

输出:
当前cuda设备是 0
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:0')

可以看到上面几个例子的在创建和转移张量时,指定了没有编号的当前cuda字符串,创建和使用的设备也是没有编号的当前cuda设备,在使用.cuda()方法时也没有指定编号或者使用了当前cuda设备。

使用torch.cuda.set_device函数可以设置任意一个可用的cuda设备为当前cuda设备,这会影响有关当前cuda设备的命令,如下所示。

复制代码
import torch

torch.cuda.set_device(1)                            # 设置当前cuda设备编号为1
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
a=torch.tensor([1, 2, 3, 4, 5], device='cuda') # 在当前cuda设备上创建一个张量
print(a)

device=torch.device('cuda')                    # 创建一个当前cuda设备

torch.cuda.set_device(2)                       # 设置当前cuda设备编号为2
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
b=torch.tensor([1, 2, 3, 4, 5], device=device) # 在当前cuda设备上创建一个张量
print(b)

torch.cuda.set_device(0)                       # 设置当前cuda设备编号为0
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
c=torch.tensor([1, 2, 3, 4, 5]).to(device)     # 将张量移动到当前cuda设备
print(c)

torch.cuda.set_device(1)                       # 设置当前cuda设备编号为1
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
d=torch.tensor([1, 2, 3, 4, 5]).to('cuda')     # 将张量移动到当前cuda设备
print(d)

torch.cuda.set_device(2)                       # 设置当前cuda设备编号为2
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
e=torch.tensor([1, 2, 3, 4, 5]).cuda()         # 将张量移动到当前cuda设备
print(e)

torch.cuda.set_device(0)                       # 设置当前cuda设备编号为0
print("当前cuda设备是", torch.cuda.current_device()) # 获取当前cuda设备
f=torch.tensor([1, 2, 3, 4, 5]).cuda(device)   # 将张量移动到当前cuda设备
print(f)

输出:
当前cuda设备是 1
tensor([1, 2, 3, 4, 5], device='cuda:1')
当前cuda设备是 2
tensor([1, 2, 3, 4, 5], device='cuda:2')
当前cuda设备是 0
tensor([1, 2, 3, 4, 5], device='cuda:0')
当前cuda设备是 1
tensor([1, 2, 3, 4, 5], device='cuda:1')
当前cuda设备是 2
tensor([1, 2, 3, 4, 5], device='cuda:2')
当前cuda设备是 0
tensor([1, 2, 3, 4, 5], device='cuda:0')

从上面的例子中,可以看出torch.cuda.set_device函数可以接受一个正整数参数作为编号,改变当前cuda设备。它还可以接受一个device对象作为参数,前提是这个device对象是有编号(即不能是当前cuda设备)。

复制代码
device =torch.device('cuda:1')
torch.cuda.set_device(device) # 这是可行的
device =torch.device('cuda')
torch.cuda.set_device(device) # 这是错误的,因为device是一个没有编号的当前cuda设备对象

输出:
ValueError: Expected a torch.device with a specified index or an integer, but got:cuda

pytorch官方不建议使用这个函数,而是直接显式指明有关cuda设备的编号,而不是依赖当前设备,如下所示。

复制代码
import torch

a=torch.tensor([1, 2, 3, 4, 5], device='cuda:1') # 在cuda1设备上创建一个张量
print(a)

device=torch.device('cuda:2')                  # 创建一个cuda2设备

b=torch.tensor([1, 2, 3, 4, 5], device=device) # 在cuda2设备上创建一个张量
print(b)

c=torch.tensor([1, 2, 3, 4, 5]).to(device)     # 将张量移动到cuda2设备
print(c)

d=torch.tensor([1, 2, 3, 4, 5]).to('cuda:0')   # 将张量移动到cuda0设备
print(d)

e=torch.tensor([1, 2, 3, 4, 5]).cuda(1)        # 将张量移动到cuda1设备
print(e)

f=torch.tensor([1, 2, 3, 4, 5]).cuda(device)   # 将张量移动到cuda2设备
print(f)

输出:
tensor([1, 2, 3, 4, 5], device='cuda:1')
tensor([1, 2, 3, 4, 5], device='cuda:2')
tensor([1, 2, 3, 4, 5], device='cuda:2')
tensor([1, 2, 3, 4, 5], device='cuda:0')
tensor([1, 2, 3, 4, 5], device='cuda:1')
tensor([1, 2, 3, 4, 5], device='cuda:2')
相关推荐
极小狐2 分钟前
极狐GitLab 如何 cherry-pick 变更?
人工智能·git·机器学习·gitlab
小宋加油啊3 分钟前
深度学习小记(包括pytorch 还有一些神经网络架构)
pytorch·深度学习·神经网络
沛沛老爹6 分钟前
从线性到非线性:简单聊聊神经网络的常见三大激活函数
人工智能·深度学习·神经网络·激活函数·relu·sigmoid·tanh
蓝莓味柯基9 分钟前
Python3:文件操作
python
0x21115 分钟前
[论文阅读]ReAct: Synergizing Reasoning and Acting in Language Models
人工智能·语言模型·自然语言处理
mucheni27 分钟前
迅为iTOP-RK3576开发板/核心板6TOPS超强算力NPU适用于ARM PC、边缘计算、个人移动互联网设备及其他多媒体产品
arm开发·人工智能·边缘计算
Jamence28 分钟前
多模态大语言模型arxiv论文略读(三十六)
人工智能·语言模型·自然语言处理
猿饵块39 分钟前
opencv--图像变换
人工智能·opencv·计算机视觉
xiaoh_71 小时前
解决视频处理中的 HEVC 解码错误:Could not find ref with POC xxx【已解决】
python·ffmpeg·音视频
LucianaiB1 小时前
【金仓数据库征文】_AI 赋能数据库运维:金仓KES的智能化未来
运维·数据库·人工智能·金仓数据库 2025 征文·数据库平替用金仓