PyTorch 1.0 中文文档:CUDA 语义

译者:片刻

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)

阅读全文/改进本文

展开阅读全文
©️2020 CSDN 皮肤主题: 黑客帝国 设计师: 上身试试 返回首页
实付0元
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、C币套餐、付费专栏及课程。

余额充值