Some question about pth2tflite or h52tflite

Viewed 24

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)
0 Answers
Related