我需要在pyTorch中创建一个固定长度的Tensor
,使其像FIFO队列一样工作。
我有以下函数来实现这个功能:
def push_to_tensor(tensor, x): tensor[:-1] = tensor[1:] tensor[-1] = x return tensor
例如,我有:
tensor = Tensor([1,2,3,4])>> tensor([ 1., 2., 3., 4.])
然后使用该函数会得到:
push_to_tensor(tensor, 5)>> tensor([ 2., 3., 4., 5.])
然而,我在想:
- pyTorch是否有原生方法可以做到这一点?
- 如果没有,是否有更聪明的方法来实现?
回答:
我实现了另一种FIFO队列:
def push_to_tensor_alternative(tensor, x): return torch.cat((tensor[1:], Tensor([x])))
功能是相同的,但随后我检查了它们在速度上的表现:
# 小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)
# 较大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)
看起来使用torch.cat
的push_to_tensor_alternative
(而不是将所有项目向左移动)速度更快。