CUDA 上的张量

最新推荐文章于 2024-07-17 10:20:14 发布

写进メ诗的结尾。

最新推荐文章于 2024-07-17 10:20:14 发布

阅读量802

点赞数

分类专栏： Pytorch 文章标签：深度学习 pytorch python 计算机视觉

本文链接：https://blog.csdn.net/weixin_48158964/article/details/127341403

版权

Pytorch 专栏收录该内容

19 篇文章 1 订阅

订阅专栏

CUDA 上的张量

torch.device() 包含一个设备类型（‘cpu’ 或 ‘cuda’）和可选的设备序号，如果设备序号不存在，则为当前设备。

>>> torch.device('cpu')
device(type='cpu')

>>> torch.device('cuda')
device(type='cuda')

>>> torch.device('cpu', 0)
device(type='cpu', index=0)

>>> torch.device('cuda', 0)
device(type='cuda', index=0)

device_1 = torch.device('cpu')
device_2 = torch.device('cuda')
print(device_1.type)
print(device_2.type)
---------
cpu
cuda

张量可以使用 to() 方法移动到任何设备（device）上。

tensor.to(‘cpu’)
tensor.to(‘cuda’)
device = torch.device(‘cuda’) --> tensor.to(device)

# 使用 torch.device() 将 tensor 移入和移出 GPU
x = torch.randn(1)
if torch.cuda.is_available():
    device = torch.device('cuda')  # CUDA 设备对象
    y = torch.ones_like(x, device=device)  # 直接在 GPU 上创建 tensor
    x = x.to(device)  # 也可以使用 x = x.to('cuda')
    z = x + y
    print(z)
    print(z.to('cpu', torch.double))  # 张量在转移设备时可以改变 dtype
---------
tensor([1.0445], device='cuda:0')
tensor([1.0445], dtype=torch.float64)