有没有办法在特定 (GPU) 设备的上下文中运行 pytorch
(无需为每个新张量指定设备,例如 .to
选项)?
类似于 tensorflow
with tf.device('/device:GPU:0'):
..
好像默认的设备是cpu(除非我做错了):
with torch.cuda.device('0'):
a = torch.zeros(1)
print(a.device)
>>> cpu
最佳答案
不幸的是,在当前的实现中,with-device
语句不能以这种方式工作,它只能用于在 cuda 设备之间切换。
您仍然必须使用 device
参数来指定使用哪个设备(或 .cuda()
将张量移动到指定的 GPU),带有像这样的术语:
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
因此要访问 cuda:1
:
cuda = torch.device('cuda')
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
并访问 cuda:2
:
cuda = torch.device('cuda')
with torch.cuda.device(2):
# allocates a tensor on GPU 2
a = torch.tensor([1., 2.], device=cuda)
但是没有 device
参数的张量仍然是 CPU 张量:
cuda = torch.device('cuda')
with torch.cuda.device(1):
# allocates a tensor on CPU
a = torch.tensor([1., 2.])
总结一下:
No - unfortunately it is in the current implementation of the
with-device
statement not possible to use in a way you described in your question.
以下是来自 documentation 的更多示例:
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)
关于python - pytorch - 使用 'with statement' 内的设备,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/52076815/