PyTorch 1.0 中文文档：CUDA 语义

Posted 2021-02-08 wizardforcel

译者：片刻

torch.cuda 用于设置和运行 CUDA 操作。它会跟踪当前选定的GPU，并且默认情况下会在该设备上创建您分配的所有 CUDA tensors。可以使用 torch.cuda.device 上下文管理器更改所选设备。

但是，一旦分配了 tensor，就可以对其进行操作而不管所选择的设备如何，结果将始终与 tensor 放在同一设备上。

默认情况下不允许跨 GPU 操作，除了 copy_() 具有类似复制功能的其他方法，例如 to() 和 cuda()。除非您启用点对点内存访问，否则任何尝试在不同设备上传播的 tensor 上启动操作都会引发错误。

下面我们用一个小例子来展示:

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)

阅读全文／改进本文