在pytorch的jit函数中可以添加列表吗?

t40tm48m  于 2023-04-21  发布在  其他
关注(0)|答案(1)|浏览(122)

我有一个创建列表并添加到列表的函数。我想将我的函数转换为Pytorch jit函数,以加快计算速度,并使用我最终将优化的参数填充列表。我不确定列表是否与Pytorch jit函数兼容,当我尝试做简单的例子时,我会出错。
比如我试着这样做

import torch

@torch.jit.script
def my_function(x):
    my_list = []
    for i in range(int(x)):
        my_list.append(i)
    return my_list

a = my_function(10)
print(a)

但我得到了这个错误

aten::append.t(t[](a!) self, t(c -> *) el) -> t[](a!):
Could not match type int to t in argument 'el': Type variable 't' previously matched to type Tensor is matched to type int.
:
  File "myscript.py", line 18
    my_list = []
    for i in range(int(x)):
        my_list.append(i)
        ~~~~~~~~~~~~~~ <--- HERE
    return my_list

这里有什么问题?我不能在PyTorch中使用列表吗?如果不能,我可以用什么其他与PyTorch兼容的可追加对象来替代?

6mw9ycah

6mw9ycah1#

你的问题是Pytorch的JIT不支持可变的列表。
您应该将列表转换为Pytorch的Tensor,然后使用torch.cat连接,以便JIT编译器知道发生了什么。

相关问题