相關閱讀
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 torchtorch.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 torcha=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')