将 pth pytorch 文件转换为 onnx 模型

Question

我正在尝试将 PyTorch 模型（包含权重的 pth 文件）转换为 onnx 文件，然后转换为 TensorFlow 模型，因为我在 TensorFlow 上工作。然后对其进行微调。到目前为止，这是我的尝试。但是我一直收到错误。enter image description here 我认为问题在于权重是用于视觉转换器的。但是我还没有弄清楚要使用什么类型的模型来转换它。我假设是 CRNN，但如果有更简单的方法，我很想知道。 PS: 我确实将 pth 文件加载到我的驱动器中。路径正确

from torch.autograd import Variable

import torch.onnx
import torchvision
import torch
import onnx
import torch.nn as nn

dummy_input = torch.randn(1, 3, 224, 224)
file_path='/content/drive/MyDrive/VitSTR/vitstr_base_patch16_224_aug.pth'

model = torchvision.models.vgg16()

model.load_state_dict(torch.load(file_path))

model.eval()

torch.onnx.export(model, dummy_input, "vitstr.onnx")

Answer 1

谢谢大家。我使用了与模型中相同的体系结构并且它起作用了。

将 pth pytorch 文件转换为 onnx 模型

Convert a pth pytorch file to an onnx model

machine-learning

deep-learning

tensorflow

pytorch

onnx