How to parallelize iterations of a tensor multiplying which requires extra indexing, or vectorize it without more memory?

user_13523104 0
Viewed 170

I have to calculate this using Pytorch. But it is very slow on GPU, actually not faster than CPU version.

# n, k           10000000         ( 1E7 )
# a        shape: (n, 100)        ( 1E9 )
# index    shape: (k,  10)        ( 1E8 )   used to select the rows of a
# w        shape: (k,  10)        ( 1E8 )
# result   shape: (k, 100)        ( 1E9 )
#
# result[i] = w[i,0]*a[index[i,0]] + ... + w[i,9]*a[index[i,9]]

for i in range(len(result)):
    result[i] = (w[i].unsqueeze(-1) * a[index[i]]).sum(0)  # shape (10, 1) (10, 100) -> (100)

To vectorize it, a simple way is removing [i]. 

result = (w.unsqueeze(-1) * a[index]).sum(1)  # shape (1E7, 10, 1) (1E7, 10, 100) -> (1E7, 100)

This version is fast. The problem is, a[index] creates a tensor of shape (k, 10, 100) i.e. (1E10), causing CUDA out of memory. I only need to access a, so the extra memory is unnecessary.

So, how to parallelize the for-loop, or how to solve the memory problem in vectorized version?

pytorch vectorization parallel-processing gpu
1 Answers

I can't really see a way around a[index]. But instead of a loop that process one line at a time, you can process chunks:

chunk = n // 10  # the largest chunk you can fit into memory
result = torch.empty((n, ), dtype=a.dtype, device=a.device)
from_ = 0
while from_ < n:
  to_ = min(from_ + chunk, n)
  result[from_:to_] = (w[from_:to_, :] * a[index[from_:to_,:]]).sum(0)
  from_ = to_
user_17144100
Related