是否可以在特定(GPU)设备的上下文中运行pytorch
(而不必为每个新张量指定设备,例如.to
选项)
类似tensorflow
with tf.device('/device:GPU:0'):
..
似乎默认设备是cpu(除非我做错了):
with torch.cuda.device('0'):
a = torch.zeros(1)
print(a.device)
>>> cpu
答案 0 :(得分:2)
不幸的是,在当前实现中,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.])
总结一下:
否-不幸的是,它在
with-device
的当前实现中 陈述不能以您在自己的描述中所述的方式使用 问题。
还有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)