时间戳

首页 > 解决方案 > 是否可以使用 pyTorch 创建 FIFO 队列？

python - 是否可以使用 pyTorch 创建 FIFO 队列？

问题描述

我需要Tensor在 pyTorch 中创建一个固定长度，其作用类似于 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.])

但是，我想知道：

标签： pythonmachine-learningqueuepytorchtensor

解决方案

我实现了另一个 FIFO 队列：

def push_to_tensor_alternative(tensor, x):
    return torch.cat((tensor[1:], Tensor([x])))

功能是相同的，但后来我检查了它们的速度性能：

# 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)

看起来像这样push_to_tensor_alternative使用torch.cat（而不是将所有项目向左移动）更快。

推荐阅读