PyTorch device與cuda.device用法介紹

2022-04-02 19:00:43

1 檢視當前的device

輸入情況：

import torch
print("Default Device : {}".format(torch.Tensor([4, 5, 6]).device))

輸出情況：

Default Device : cpu

輸入情況

device = torch.Tensor([1, 2, 3], device="cpu:0").device
print("Device Type: {}".format(device))

輸出情況

Device Type: cpu

輸入情況

gpu = torch.device("cuda:0")
print("GPU Device:【{}：{}】".format(gpu.type, gpu.index))

輸出情況

GPU Device:【cuda：0】

輸入情況

print("Total GPU Count :{}".format(torch.cuda.device_count()))
print("Total CPU Count :{}".format(torch.cuda.os.cpu_count()))

輸出情況

Total GPU Count :1
Total CPU Count :8

輸入情況

data = torch.Tensor([[1, 4, 7], [3, 6, 9], [2, 5, 8]])
print(data.shape)

輸出情況

torch.Size([3, 3])

輸入情況：

data_gpu = data.to(torch.device("cuda:0"))
print(data_gpu.device)

輸出情況：

cuda:0

輸入情況：

data_gpu2 = data.cuda(torch.device("cuda:0"))
# 如果只有一塊gpu的話  直接寫成這樣：data_gpu2 = data.cuda()
print(data_gpu2.device)

輸出情況：

cuda:0

到此這篇關於PyTorch device與cuda.device用法的文章就介紹到這了,更多相關PyTorch device使用內容請搜尋it145.com以前的文章或繼續瀏覽下面的相關文章希望大家以後多多支援it145.com！