pytorch学习（十六）to(device) 和.cuda()的用法

最新推荐文章于 2025-03-21 21:12:44 发布

TEn%

最新推荐文章于 2025-03-21 21:12:44 发布

阅读量1w

点赞数 3

分类专栏： Pytorch系列学习 Python 文章标签： cuda gpu 深度学习

本文链接：https://blog.csdn.net/qq_37297763/article/details/118637442

版权

Python 同时被 2 个专栏收录

21 篇文章

订阅专栏

Pytorch系列学习

18 篇文章

订阅专栏

这篇博客介绍了在PyTorch中如何将张量转移到GPU设备，包括直接使用`cuda()`函数和通过`to(device)`方法。这两种方法都能实现张量在不同GPU之间的分配和转移，并且在指定设备上下文中操作。示例代码展示了如何在GPU0和GPU2上创建和操作张量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

看pytorch官方给的代码，to(device)的时候可以指定GPU设备号。二者没啥太大区别，都可以用。

cuda = torch.device('cuda')     # Default CUDA device
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2')  # GPU 2 (these are 0-indexed)

x = torch.tensor([1., 2.], device=cuda0)
# x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
# y.device is device(type='cuda', index=0)

with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.tensor([1., 2.], device=cuda)

    # transfers a tensor from CPU to GPU 1
    b = torch.tensor([1., 2.]).cuda()
    # a.device and b.device are device(type='cuda', index=1)

    # You can also use ``Tensor.to`` to transfer a tensor:
    b2 = torch.tensor([1., 2.]).to(device=cuda)
    # b.device and b2.device are device(type='cuda', index=1)

    c = a + b
    # c.device is device(type='cuda', index=1)

    z = x + y
    # z.device is device(type='cuda', index=0)

    # even within a context, you can specify the device
    # (or give a GPU index to the .cuda call)
    d = torch.randn(2, device=cuda2)
    e = torch.randn(2).to(cuda2)
    f = torch.randn(2).cuda(cuda2)
    # d.device, e.device, and f.device are all device(type='cuda', index=2)