Sometimes, current operations tend to hardly meet our needs as they are not ‘vectorized enough’.
For example, here I have to
tensor.eq(float) to get each
cur_sp_mask. So I have to use
for-loop to enumerate each blob.
import torch bz = 2 a = (torch.rand(bz, 3, 4)*4).int() index = torch.tensor([[0, 1, 0, 2],[0, 1, 2, 3], [1, 2, 2, 3]]) a = a.view(bz, -1) index_vec = index.view(-1) for i in range(bz): max_index = torch.max(index_vec) a_vec = a[i] rec = torch.zeros_like(a_vec.float()) # I have to enumerate each blob for j in range(max_index + 1): cur_sp_mask = index_vec.eq(j) n_pixels = torch.sum(cur_sp_mask).float() rec[cur_sp_mask] = a_vec[cur_sp_mask] + j * n_pixels rec = rec.view(3, 4) print('a[i]') print(a_vec.view(3,4)) print('index') print(index.view(3,4)) print('result') print(rec) print('------')
Another example, I want to create a list of tensors, each of which can be created using
# Something like this: # import torch x= torch.tensor([20, 53, 98]).long() y= torch.tensor([110, 120, 132]).long() # now I have to a =  for i in range(3): a.append(torch.arange(x[i], y[i])) # However, I hope something like this can be # more efficient than `for-loop`. a = torch.arange(x, y) # a is a list full of tensors
torch.arange does not support “batching” operation.
It is definetely huge work to make most operations powerful enough to support so advanced “batching”. However, maybe it is possible:
import torch x= torch.tensor([20, 53, 98]).long() y= torch.tensor([110, 120, 132]).long() warp_x = warp_op(x) warp_y = warp_op(y) # When tensors are warpped, # then they can be more efficient executed than `for-loop`. a = torch.arange(warp_x , warp_y) # a is a list full of tensors