pytorch - использовать устройство внутри 'with Statement' - PullRequest
Новая
       8

pytorch - использовать устройство внутри 'with Statement'

0 голосов
/ 29 августа 2018

Есть ли способ запустить pytorch в контексте определенного устройства (GPU) (без необходимости указывать устройство для каждого нового тензора, например, параметр .to)?

Что-то вроде эквивалента tensorflow with tf.device('/device:GPU:0'): ..

Похоже, что устройством по умолчанию является процессор (если я не делаю это неправильно): 

with torch.cuda.device('0'):
   a = torch.zeros(1)
   print(a.device)

>>> cpu

1 Ответ

0 голосов
/ 29 августа 2018

К сожалению, в текущей реализации оператор with-device не работает таким образом, его можно просто использовать для переключения между устройствами cuda. ​​

Вам все равно придется использовать параметр device, чтобы указать, какое устройство используется (или .cuda(), чтобы переместить тензор в указанный графический процессор) с такой терминологией, как: 

# 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 все равно будут тензорами процессора: 

cuda = torch.device('cuda')

with torch.cuda.device(1):
    # allocates a tensor on CPU
    a = torch.tensor([1., 2.])

Подводя итог:

Нет - к сожалению, это в текущей реализации with-device Заявление невозможно использовать так, как вы описали в своем вопрос.

Вот еще несколько примеров из документации : 

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)
...