标签:from 设备 use strong 1.0 device tle 一个 运行
译者:片刻
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)
标签:from 设备 use strong 1.0 device tle 一个 运行
原文地址:https://www.cnblogs.com/wizardforcel/p/10358928.html