Могут ли массивы numpy работать в графических процессорах?

Question

Я использую PyTorch. У меня есть следующий код:

import numpy as np
import torch

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6], [1, 2, 3, 4]])
X = torch.DoubleTensor(X).cuda()

X_split = np.array_split(X.numpy(), 
                         indices_or_sections = 2, 
                         axis = 0)
X_split

, но я получаю эту ошибку:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-121-870b5d3f67b6> in <module>()
----> 1 X_prime_class_split = np.array_split(X_prime_class.numpy(), 
      2                                      indices_or_sections = 2,
      3                                      axis = 0)
      4 X_prime_class_split

TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

Сообщение об ошибке ясное, и я знаю, как исправить эту ошибку, просто включив .cpu() , ie. X_prime_class.cpu().numpy(). Мне просто любопытно узнать, подтверждает ли это, что массивы numpy не могут работать в графических процессорах / Cuda?

Answer 1

Нет, вы не можете запускать numpy функции на массивах GPU. PyTorch повторно реализует большую часть функций numpy для тензоров PyTorch. Например, torch.chunk работает аналогично np.array_split, поэтому вы можете сделать следующее:

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6], [1, 2, 3, 4]])
X = torch.DoubleTensor(X).cuda()
X_split = torch.chunk(X, chunks=2, dim=0)

, который разбивает X на несколько тензоров, не перемещая X с Графический процессор.