I'm wondering if it's possible to do something in PyTorch...
Basically, what I'm trying to do is something like torch.narrow(), but where the start value is a tensor of indices.
Suppose I had the following tensors:
start_indices = [3, 1, 2, ...] # shape: (M,)
dataset = [ 3, 5, 3, 4, 8, 0, 1, 3, 9, 7, 2, 7, 3, 7, 6, 0, 2, 3, 0, 2, 5, ...
] # shape: (M, N) (here, N = 7, just for demo)In my use case, M ~ 10^3 and N > 10^6
I'm wondering if it's possible to do something like this:
torch.narrow( input=dataset, dim=1, start=start_indices, length=4
)that returns a "narrowed" version of the original dataset, where the start of each row is selected by its corresponding start index. To complete the example I gave above, the result would look like:
result = [ 4, 8, 0, 1, 9, 7, 2, 7, 2, 3, 0, 2, ....
] # shape: (M, 4)I realize that depending on the start indices and length, you could have some lines going out of bounds of the original tensor -> I'm fine with receiving an error in this case.
Right now, I'm currently using a for-loop to do this, but it becomes far too slow as M becomes large. I'm hoping there's a way to do this that takes advantage of data parallelism.
Any ideas? I thought of using torch.index_select(), but it only selects within rows based off of single indices. I want to select within rows by single indices, but then extend those selections out by some fixed length. In that way, what I'm looking for is kind of a "combination" of torch.narrow() and torch.index_select().
Thanks in advance!
Related questions 8 Filling torch tensor with zeros after certain index 0 Pytorch - select region of a tensor using torch function 7 Index a torch tensor with an array Related questions 8 Filling torch tensor with zeros after certain index 0 Pytorch - select region of a tensor using torch function 7 Index a torch tensor with an array 0 How to convert a matrix of torch.tensor to a larger tensor? 3 How can I get argmaxed torch tensor excluding certain index? 2 Efficiently filling torch.Tensor at equal index positions 1 Appling sliding window to torch.tensor and adjusting tensor initial size 4 Split a torch tensor using a same-sized tensor of indices 1 Reduce torch tensor 2 How to use indices from torch.min to index multi-dimensional tensor Load 7 more related questions Show fewer related questions Reset to default