将pth pytorch文件转换为onnx模型

Question

我试图将PyTorch模型(包含权重的第pth文件)转换为onnx文件，然后转换为TensorFlow模型，因为我在TensorFlow上工作。然后对它进行微调。这是我目前的尝试。然而，我总是出错。在这里输入图像描述我认为问题在于重量是用于视觉转换器的。但是我还不知道用什么模型来转换它。我假设是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

谢谢大家。我使用了与模型中相同的架构，并且它有效。