JIT 保存 pytorch 模型时出现未知类型注释错误

Question

当 JIT 保存具有许多自定义类的复杂 pytorch 模型的“model.pt”时，我遇到了 pytorch 不知道这些自定义类之一的类型注释的错误。换句话说，以下代码（对原始代码进行了彻底总结）在第七行失败：

import torch

from gan import Generator

from gan.blocks import SpadeBlock

generator = Generator()

generator.load_weights("path/to/weigts")

jitted = torch.jit.script(generator)

torch.jit.save(jitted, "model.pt")

错误：

Tracebck (most recent call last):

  File "pth2onnx.py", line 72, in <module>

    to_torch_jit(generator)

  File "pth2onnx.py", line 24, in to_torch_jit

    jitted = torch.jit.script(generator)

  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 1516, in script

    return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)

  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 310, in create_script_module

    concrete_type = concrete_type_store.get_or_create_concrete_type(nn_module)

  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 269, in get_or_create_concrete_type

    concrete_type_builder = infer_concrete_type_builder(nn_module)

  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 138, in infer_concrete_type_builder

    sub_concrete_type = concrete_type_store.get_or_create_concrete_type(item)

  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 269, in get_or_create_concrete_type

它抱怨的类型确实是我们自己编写并在加载的Generator. 我将不胜感激有关可能导致此问题的原因或如何调查此问题的指示！

我尝试了以下方法：

• 在调用 torch.jit.script 的脚本中显式导入 SpadeBlock

• 确保它继承自 nn.Module （生成器也是如此）

• 使用 pip install --user -e 确保安装了 gan 软件包

有任何想法吗？提前致谢！

Answer 1

问题原来是我使用的类变量名称被破坏了。例子：class&nbsp;Generator(nn.Module): &nbsp;&nbsp;&nbsp;&nbsp;__main:&nbsp;nn.Module两个前导下划线就是原因。将它们更改为单个下划线或无下划线。解决问题。class&nbsp;Generator(nn.Module): &nbsp;&nbsp;&nbsp;&nbsp;main:&nbsp;nn.Module