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

Question

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

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")

错误：

Traceback (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
    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
    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 126, in infer_concrete_type_builder
    attr_type = infer_type(name, item)
  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/_recursive.py", line 99, in infer_type
    attr_type = torch.jit.annotations.ann_to_type(class_annotations[name], _jit_internal.fake_range())
  File "/home/a.nieuwland/.conda/envs/python3.6/lib/python3.6/site-packages/torch/jit/annotations.py", line 303, in ann_to_type
    raise ValueError("Unknown type annotation: '{}'".format(ann))
ValueError: Unknown type annotation: '<class 'gan.blocks.SpadeBlock'>'

它抱怨的类型确实是我们自己编程并在加载的Generator中使用的class。我将不胜感激关于可能导致此问题或如何调查此问题的指示！

我尝试了以下方法：

在调用 torch.jit.script
确保它继承自 nn.Module（与生成器一样）
确保安装了 gan 包，使用 pip install --user -e

有什么想法吗？提前致谢！

Answer 1

问题原来是我使用的 class 变量名被破坏了。示例：

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

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

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

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

Getting Unknown type annotation error when JIT saving pytorch model

python

jit

pytorch