我有一个创建列表并添加到列表的函数。我想将我的函数转换为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兼容的可追加对象来替代?
1条答案
按热度按时间6mw9ycah1#
你的问题是Pytorch的JIT不支持可变的列表。
您应该将列表转换为Pytorch的Tensor,然后使用
torch.cat
连接,以便JIT编译器知道发生了什么。