Ошибка при преобразовании модели PyTorch в формат ONNX ⇐ Python

[Anonymous]

Я хочу преобразовать https://github.com/TencentARC/GFPGAN в формат ONNX для использования в приложении Android. После выполнения процедуры отображается ошибка (AttributeError: объект «dict» не имеет атрибута «модули»).
Ниже приведен мой код в Colab:

Код: Выделить всё

!pip install onnx

!pip install torch torchvision onnx
import urllib.request

url = "https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth"
urllib.request.urlretrieve(url, "my_model.pth")

Код: Выделить всё

import torch
import numpy as np

model = torch.load("my_model.pth")

input_shape = (1, 3, 224, 224)
example_input = torch.rand(input_shape)

Код: Выделить всё

import urllib.request
import torch

output_path = "my_model.onnx"
torch.onnx.export(model, input_shape, output_path, opset_version=11)

# Load and check the ONNX model for correctness
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)

Здесь в последнем фрагменте кода отображается следующая ошибка:

Код: Выделить всё

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
 in ()
3
4 output_path = "my_model.onnx"
----> 5 torch.onnx.export(model, input_shape, output_path, opset_version=11)
6
7 # Load and check the ONNX model for correctness

5 frames
/usr/local/lib/python3.9/dist-packages/torch/onnx/utils.py in disable_apex_o2_state_dict_hook(model)
137     if not isinstance(model, torch.jit.ScriptFunction):
138         model_hooks = {}  # type: ignore[var-annotated]
--> 139         for module in model.modules():
140             for key, hook in module._state_dict_hooks.items():
141                 if type(hook).__name__ == "O2StateDictHook":

AttributeError: 'dict' object has no attribute 'modules'

Я попробовал преобразовать модель в ONNX и импортировать библиотеки. Но проблема возникает при экспорте torch.onnx.export.

Подробнее здесь: https://stackoverflow.com/questions/760 ... nnx-format