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

编程问答 2022-08-11

问题描述

当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软件包已安装

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

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）

jitpythonpytorch