究极转换：pth转换kmodel

究极折磨王

提示：这个是自己的笔记，后悔当时为了偷懒，而先学习pytorch后学tensorflow了

---

---

前言

提示：这篇文章引用了大家的思路进行设计，找了网上好多篇文章最后的总结。

---

提示：以下是本篇文章正文内容，下面案例可供参考

一、pth文件转换onnx

这步没啥可说的，直接上传代码：

import torch
import torch.onnx
from tinynet import tinynet
from conf import settings
import os

def pth_to_onnx(input, checkpoint, onnx_path, input_names=['input'], output_names=['output'], device='cpu'):
    if not onnx_path.endswith('.onnx'):
        print('Warning! The onnx model name is not correct,\
              please give a name that ends with \'.onnx\'!')
        return 0

    model = tinynet()
    model.load_state_dict(torch.load(checkpoint))
    model.eval()
    # model.to(device)
    
    torch.onnx.export(model, input, onnx_path, verbose=True, input_names=input_names, output_names=output_names)
    print("Exporting .pth model to onnx model has been successful!")

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES']='2'
    checkpoint = './xiaoan.pth'
    onnx_path = './xiaoan.onnx'
    input = torch.randn(-1, 3, 32, 32)  #这里用的是自己的模型，数据需要根据自己的模型进行更改。
    # device = torch.device("cuda:2" if torch.cuda.is_available() else 'cpu')
    pth_to_onnx(input, checkpoint, onnx_path)

二、onnx转换pb

这步可是老折磨王了，废话不多说，直接上代码：
这步一定一定要注意维度的转换。

pip install -e.
pip install tensorflow-addons


onnx-tf convert -i /xiaoan.onnx -o /xiaoan.pb

代码转换
import onnx
from onnx_tf.backend import prepare
import os

def onnx2pb(onnx_input_path, pb_output_path):
    onnx_model = onnx.load(onnx_input_path)  # load onnx model
    tf_exp = prepare(onnx_model)  # prepare tf representation
    tf_exp.export_graph(pb_output_path)  # export the model


if __name__ == "__main__":
    os.makedirs("tensorflow", exist_ok=True)
    onnx_input_path = './onnx/xiaoan.onnx'
    pb_output_path = './onnx/xiaoan.pb'

    onnx2pb(onnx_input_path, pb_output_path)

pb转换tflite

这步就可以了，毕竟tensorflow内部的程序，就比较好转换。

import tensorflow as tf
 
in_path=r'D:\zhuanhuan\xiaoan.pb'
out_path=r'D:\zhuanhuan\xiaoan.tflite'
 
input_tensor_name=['train']
input_tensor_shape={'xiaoan':[-1,32,32,3]}
 
class_tensor_name=['heisibing','heixingbing','xiubing','jiankang']
 
 
convertr=tf.lite.TFLiteConverter.from_frozen_graph(in_path,input_arrays=input_tensor_name
                                                   ,output_arrays=class_tensor_name
                                                   ,input_shapes=input_tensor_shape)
 
# convertr=tf.lite.TFLiteConverter.from_saved_model(saved_model_dir=in_path,input_arrays=[input_tensor_name],output_arrays=[class_tensor_name])
tflite_model=convertr.convert()
 
 
with open(out_path,'wb') as f:
    f.write(tflite_model)

tflite转换kmodel

下载工具NNCase工具箱
创建一个文件夹,将下载好的ncc.exe,.tflite文件,训练模型时的图片(5张即可)放入文件夹内
然后直接跑代码：

ncc compile xiaoan.tflite xiaoan.kmodel -i tflite -o kmodel --dataset 'D/zhuanhuan/train'