I have got a pytorch pretrained model and want to convert pth to tflite(including resnet50 or mobilevit).
Firstly, i have concerted the model in pth formart to ONNX format which will not find so many transpose layers.
Then, I convert the ONNX model into tensorflow which using onnx-tf.
Nextly, I convert the tensorflow model into Tensorflow Lite and find some many unexpected traspose layers which will make some bad influnce on deploying the model on MCU.
Below is a minimised example of what I am doing.
import torch
import onnx
from onnx_tf.backend import prepare
import tensorflow as tf
import torchvision
import os
import time
import timm
from pytorch_pretrained_vit import ViT
def pth_to_onnx(output_path):
torch_model = timm.create_model('resnet50', pretrained=True)
x = torch.randn(1, 3, 224, 224)
export_onnx_file = output_path
torch.onnx.export(torch_model,
x,
export_onnx_file,
opset_version=11,
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch_size"},
"output": {0: "batch_size"}}
)
def onnx_to_pb(output_path):
model = onnx.load(output_path)
tf_rep = prepare(model)
tf_rep.export_graph('resnet50')
if __name__=='__main__':
output_path = "resnet50.onnx"
pth_to_onnx(output_path)
onnx_to_pb(output_path)
saved_model_dir = 'resnet50'
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) # path to the SavedModel directory
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
# Save the model.
with open('resnet50.tflite', 'wb') as f:
f.write(tflite_model)