pytorch torch.jit.trace returns 函数而不是 torch.jit.ScriptModule

Question

我需要在 C++ 中运行预训练的 pytorch nn 模型（在 python 中训练）来进行预测。

为此，我按照此处给出的有关如何在 C++ 中加载 pytorch 模型的说明进行操作：https://pytorch.org/tutorials/advanced/cpp_export.html

但是当我尝试按照教程第一步中所述通过跟踪获取 torch.jit.ScriptModule 时：

    traced_script_module =
        torch.jit.trace(model, (input_tensor_1, input_tensor_2))

不是返回 torch.jit.ScriptModule，而是 returns 一个函数：

    print(type(traced_script_module))
    <type 'function'>

其中，当我运行:

    traced_script_module.save("model.pt")

然后会导致以下错误：

Traceback (most recent call last):
  File "serialize_model.py", line 60, in <module>
    traced_script_module.save("model.pt")
AttributeError: 'function' object has no attribute 'save'

对我做错了什么有什么想法吗？

Answer 1

感谢您的提问Jatentaki。我在 Python 中使用 PyTorch 0.4，当我更新到 1.0 时它起作用了。

pytorch torch.jit.trace returns 函数而不是 torch.jit.ScriptModule

pytorch torch.jit.trace returns function instead of torch.jit.ScriptModule

c++

python

pytorch

libtorch