PyTorch——device与cuda.device用法-Toy模板网

这篇具有很好参考价值的文章主要介绍了PyTorch——device与cuda.device用法。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

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

文章来源地址https://www.toymoban.com/news/detail-404890.html

到了这里，关于PyTorch——device与cuda.device用法的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！