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

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

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

错误：

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 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中使用的类。对于可能导致此问题或如何进行调查的指示，我将不胜感激！

我尝试了以下操作：

• 在调用torch.jit.script的脚本中明确导入SpadeBlock
• 确保它从nn.Module继承（与Generator一样）
• 使用pip install --user -e
确保gan软件包已安装

有什么想法吗？预先感谢！

解决方法

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

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

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

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