首页 > 解决方案 > 将预训练的 Torch 模型转换为 Tensorflow

问题描述

我想使用 Torch Image Models Library 中的预训练模型并将其转换为 tensorflow 模型,但我遇到了问题。我一直在关注官方的 onnx教程

from torchsummary import summary
import torch
from torch.autograd import Variable
import timm
import onnx
from onnx_tf.backend import prepare
from tensorflow.keras.models import load_model

model = timm.create_model('vit_tiny_r_s16_p8_224', pretrained=True)

dummy_var = Variable(torch.randn(1, 3, 224, 224))
torch.onnx.export(model, dummy_var, 'torch_model', opset_version=11 )

#load at onnx model
model = onnx.load('torch_model')
tf_rep = prepare(model)

tf_rep.export_graph('tf_model')

tf_model = load_model('tf_model')

WARNING:tensorflow:SavedModel 在加载 Keras 模型时检测到在 TF 2.5 之前保存。请确保您使用 model.save() 或 tf.keras.models.save_model() 保存模型,而不是tf.saved_model.save()。为了确认,在 SavedModel 目录中应该有一个名为“keras_metadata.pb”的文件。

然后收不到摘要

tf_model.summary()

AttributeError: '_UserObject' 对象没有属性 'summary' AttributeError Traceback (最近一次调用最后) /tmp/ipykernel_27669/3932266756.py in ----> 1 tf_model.summary()

AttributeError:“_UserObject”对象没有属性“摘要”

这是因为以下

print(tf_model)

<tensorflow.python.saved_model.load.Loader._recreate_base_user_object.._UserObject at 0x7f24ad939b20>

标签: pythontensorflowkeraspytorchonnx

解决方案


推荐阅读