Most Efficient Way to Get Ordered Slice of Array as Tensor

1

Suppose we have an array:

A = torch.arange(start = 1, end = 101)

And we want to take a rolling window over the array to produce matrix B such that:

B=([1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
      [2, 3, 4, ...],
      [3, 4, 5, ...],
      [...],
      [91, 92, 93, 94, ..., 99, 100])

I know we could run a loop through to make it, but is there a faster, parallelized way in Pytorch to render this?

2 
A = torch.arange(start = 1, end = 101)
B = (A.reshape(-1, 10))

Any idea?

3

Hi @thecho7 ,

That is not quite what I am looking for. Here is an example via loop:

A=torch.arange(start=1, end=101)

def stack_sequence(sequence, window_size):
    thread_range=sequence.size(-1)-window_size+1
    out=torch.empty(0,window_size)
    for i in range(thread_range):
        out=torch.cat([out,sequence[i:i+window_size].unsqueeze(0)])

    return out

B = stack_sequence(A,10)
print(B)

What I am looking for is a more efficient solution, if one exists.

4 
torch.stack([A[x:x+10] for x in range(91)])
1 Like
5

Thanks @anantguptadbl . That works much faster than the loop.

import torch
import time


def stack_sequence(sequence, window_size):
    thread_range=sequence.size(-1)-window_size+1
    out=torch.empty(0,window_size)
    for i in range(thread_range):
        out=torch.cat([out,sequence[i:i+window_size].unsqueeze(0)])

    return out

def stack_sequence2(sequence, window_size):
    return torch.stack([sequence[x:x+window_size] for x in range(sequence.size()[-1]-window_size+1)])



A=torch.arange(start=1, end=100001)
now=time.time()
stack_sequence(A, 10)
print(time.time()-now)

now=time.time()
stack_sequence2(A,10)
print(time.time()-now)

Results:

32.05299949645996
0.2589995861053467