Pytorch学习之cuda

最新推荐文章于 2025-03-13 08:55:24 发布

笃静悟初

最新推荐文章于 2025-03-13 08:55:24 发布

阅读量1.3w

点赞数 1

分类专栏： pytorch学习文章标签： cuda

pytorch学习专栏收录该内容

7 篇文章

订阅专栏

本文介绍了如何在PyTorch中使用CUDA进行GPU运算。torch.cuda模块用于管理和执行CUDA操作，分配的CUDA tensors会根据默认设备创建。运算遵循设备就地原则，跨GPU操作需借助copy_(), to()或cuda()函数。示例展示了CUDA的使用。" 114622257,10293654,Spring任务调度详解：TaskExecutor与TaskScheduler,"['Java', 'Spring框架', '任务调度', '并发处理', '异步编程']

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

CUDA 语义

torch.cuda用来设置和运行CUDA操作。它会记录下当前所选的GPU,这样一来，你分配的所有CUDA tensors默认就会创建在默认的设备（cpu或者某个gpu）上。当然默认的设备也可以通过torch.cuda.device内容管理器来进行改变。

然而，一旦分配了tensor之后，你不需要考虑你最初所选的设备（cpu或者某个gpu），运算最终的结果一定会跟你的tensor所选的设备是一样的。

跨GPU的运算默认是不允许的，但是可以通过copy_()和类似copy功能的方法to()和cuda()来进行实现跨GPU操作（其实还是在一个GPU上，就是把另一个GPU或者cpu的tensor给拷贝过来而已）。除非你能够实现对等的内存访问，否则任何试图在跨不同设备对tensor进行操作会触发错误。

Below you can find a small example showcasing this:

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)

翻译自：https://pytorch.org/docs/stable/notes/cuda.html