JOHNKYON
JOHNKYON

Reputation: 83

Can we use pytorch scatter_ on GPU

I'm trying to do one hot encoding on some data with pyTorch on GPU mode, however, it keeps giving me an exception. Can anybody help me?

Here's one example:

def char_OneHotEncoding(x):
    coded = torch.zeros(x.shape[0], x.shape[1], 101)
    for i in range(x.shape[1]):
        coded[:,i] = scatter(x[:,i])
    return coded


def scatter(x):
    return torch.zeros(x.shape[0], 101).scatter_(1, x.view(-1,1), 1)

So if I give it an tensor on GPU, it shows like this:

x_train = [[ 0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0],
       [14, 13, 83, 18, 14],
       [ 0,  0,  0,  0,  0]]
print(char_OneHotEncoding(torch.tensor(x_train, dtype=torch.long).cuda()).shape)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-62-95c0c4ade406> in <module>()
      4        [14, 13, 83, 18, 14],
      5        [ 0,  0,  0,  0,  0]]
----> 6 print(char_OneHotEncoding(torch.tensor(x_train, dtype=torch.long).cuda()).shape)
      7 x_train[:5, maxlen:maxlen+5]

<ipython-input-53-055f1bf71306> in char_OneHotEncoding(x)
      2     coded = torch.zeros(x.shape[0], x.shape[1], 101)
      3     for i in range(x.shape[1]):
----> 4         coded[:,i] = scatter(x[:,i])
      5     return coded
      6 

<ipython-input-53-055f1bf71306> in scatter(x)
      7 
      8 def scatter(x):
----> 9     return torch.zeros(x.shape[0], 101).scatter_(1, x.view(-1,1), 1)

RuntimeError: Expected object of backend CPU but got backend CUDA for argument #3 'index'

BTW, if we simply remove the .cuda() here, everything goes one well

print(char_OneHotEncoding(torch.tensor(x_train, dtype=torch.long)).shape)
torch.Size([5, 5, 101])

Upvotes: 3

Views: 3288

Answers (1)

Jatentaki
Jatentaki

Reputation: 13113

Yes, it is possible. You have to pay attention that all tensors are on GPU. In particular, by default, constructors like torch.zeros allocate on CPU, which will lead to this kind of mismatches. Your code can be fixed by constructing with device=x.device, as below

import torch

 def char_OneHotEncoding(x):
     coded = torch.zeros(x.shape[0], x.shape[1], 101, device=x.device)
     for i in range(x.shape[1]):
         coded[:,i] = scatter(x[:,i])
     return coded


 def scatter(x):
     return torch.zeros(x.shape[0], 101, device=x.device).scatter_(1, x.view(-1,1), 1)

 x_train = torch.tensor([
     [ 0,  0,  0,  0,  0],
     [ 0,  0,  0,  0,  0],
     [ 0,  0,  0,  0,  0],
     [14, 13, 83, 18, 14],
     [ 0,  0,  0,  0,  0]
 ], dtype=torch.long, device='cuda')

 print(char_OneHotEncoding(x_train).shape)

Another alternative are constructors called xxx_like, for instance zeros_like, though in this case, since you need different shapes than x, I found device=x.device more readable.

Upvotes: 2

Related Questions