Reputation: 2121
I need to create a fixed length Tensor
in pyTorch that acts like a FIFO queue.
I have this fuction to do it:
def push_to_tensor(tensor, x):
tensor[:-1] = tensor[1:]
tensor[-1] = x
return tensor
For example, I have:
tensor = Tensor([1,2,3,4])
>> tensor([ 1., 2., 3., 4.])
then using the function will give:
push_to_tensor(tensor, 5)
>> tensor([ 2., 3., 4., 5.])
However, I was wondering:
Upvotes: 5
Views: 6011
Reputation: 15119
A more generic version of @Bruno_Lubascher's answer, with control over the deque size, support of batched insertion, and control over the push dimension:
def push_to_deque(deque, x, deque_size=None, dim=0):
"""Handling `deque` tensor as a (set of) deque/FIFO, push the content of `x` into it."""
if deque_size is None:
deque_size = deque.shape[dim]
deque_dims = deque.dim()
input_size = x.shape[dim]
dims_right = deque_dims - dim - 1
deque_slicing = (
(slice(None),) * dim
+ (
slice(
input_size - deque_size
if input_size < deque_size
else deque.shape[dim],
None,
),
)
+ (slice(None),) * dims_right
)
input_slicing = (
(slice(None),) * dim + (slice(-deque_size, None),) + (slice(None),) * dims_right
)
deque = torch.cat((deque[deque_slicing], x[input_slicing]), dim=dim)
return deque
Examples:
>>> # Consider batched deques containing vectors of shape (2,):
>>> batch_size, vector_size = 1, 2
>>> deque_size = 4
>>> # Initialize the empty deques:
>>> deques = torch.empty((batch_size, 0, vector_size))
>>> # Push at once more vectors than the batched FIFOs can contain:
>>> vals = torch.arange(10).view((batch_size, 5, vector_size))
>>> deque = push_to_deque(deque, vals, deque_size=deque_size, dim=1)
>>> deque
tensor([[[2., 3.],
[4., 5.],
[6., 7.],
[8., 9.]]])
>>> # Push some more:
>>> vals = torch.arange(10, 20).view((batch_size, 5, vector_size))
>>> deque = push_to_deque(deque, vals, deque_size=deque_size, dim=1)
>>> deque
tensor([[[12., 13.],
[14., 15.],
[16., 17.],
[18., 19.]]])
>>> vals = torch.arange(20, 24).view((batch_size, 2, vector_size))
>>> deque = push_to_deque(deque, vals, deque_size=deque_size, dim=1)
>>> deque
tensor([[[16., 17.],
[18., 19.],
[20., 21.],
[22., 23.]]])
>>> # Verify the method can also handle oversized FIFOs:
>>> deque = torch.zeros(batch_size, 10, vector_size)
>>> vals = torch.arange(4).view((batch_size, 2, vector_size))
>>> deque = push_to_deque(deque, vals, deque_size=deque_size, dim=1)
>>> deque
tensor([[[0., 0.],
[0., 0.],
[0., 1.],
[2., 3.]]])
Upvotes: 1
Reputation: 416
Maybe a little late but I found another way to do this and save some time.
In my case, I needed a similar FIFO structure but I only needed to actually parse the
FIFO tensor once every N iterations. i.e. I needed a FIFO structure to hold n
integers, and every n
iterations I needed to parse that tensor thourgh my model. I found it is way faster to implement a collections.deque
structure and cast that deque to a tensor torch.
import time
import torch
from collections import deque
length = 5000
que = deque([0]*200)
ten = torch.tensor(que)
s = time.time()
for i in range(length):
for j in range(200):
que.pop()
que.appendleft(j*10)
torch.tensor(que)
# after some appending/popping elements, cast to tensor
print("finish deque:", time.time()-s)
s = time.time()
for i in range(length):
for j in range(200):
newelem = torch.tensor([j*10])
ten = torch.cat((ten[1:], newelem))
#using tensor as FIFO, no need to cast to tensor
print("finish tensor:", time.time()-s)
the results are the following:
finish deque: 0.15857529640197754
finish tensor: 9.483643531799316
I also noticed that when using a deque and always casting to a torch.tensor instead
of using push_alternative
it can give you a ~20% boost in time.
s = time.time()
for j in range(length):
que.pop()
que.appendleft(j*10)
torch.tensor(que)
print("finish queue:", time.time()-s)
s = time.time()
for j in range(length):
newelem = torch.tensor([j*10])
ten = torch.cat((ten[1:], newelem))
print("finish tensor:", time.time()-s)
results:
finish queue: 8.422480821609497
finish tensor: 11.169137477874756
Upvotes: 2
Reputation: 2121
I implemented another FIFO queue:
def push_to_tensor_alternative(tensor, x):
return torch.cat((tensor[1:], Tensor([x])))
The functionality is the same, but then I checked their performance in speed:
# Small Tensor
tensor = Tensor([1,2,3,4])
%timeit push_to_tensor(tensor, 5)
>> 30.9 µs ± 1.26 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit push_to_tensor_alternative(tensor, 5)
>> 22.1 µs ± 2.25 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# Larger Tensor
tensor = torch.arange(10000)
%timeit push_to_tensor(tensor, 5)
>> 57.7 µs ± 4.88 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit push_to_tensor_alternative(tensor, 5)
>> 28.9 µs ± 570 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
Seems like this push_to_tensor_alternative
which uses torch.cat
(instead of shifting all items to the left) is faster.
Upvotes: 9