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

Question

当 JIT 保存具有许多自定义类的复杂 pytorch 模型的“model.pt”时，我遇到了 pytorch 不知道这些自定义类之一的类型注释的错误。换句话说，以下代码（对原始代码进行了彻底总结）在第七行失败：import torchfrom gan import Generatorfrom gan.blocks import SpadeBlockgenerator = 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 Generator(nn.Module):
    __main: nn.Module

两个前导下划线就是原因。将它们更改为单个下划线或无下划线。解决问题。

class Generator(nn.Module):
    main: nn.Module