YOLOV3——PyTorch训练TensorFlowLite部署模型转换

Posted 2021-03-10 thunder-wu

• 主要思路 

          将训练好的.pt文件转换为keras的.h5文件，再将.h5文件转换为.tflite文件。

 

• 步骤：

       1.环境：PyTorch1.0以上（其他版本也可以，主要是适配下面的开源代码）

       .pt转weights下载后在文件内添加pt_to_weights.py。

from models import *

model = Darknet("cfg/yolov3-obj.cfg")#你用于训练的cfg文件
load_darknet_weights(model,"weights/latest.pt")
save_weights(model,path=‘weights/latest.weights‘,cutoff=-1)

      2.环境：Tensorflow2.0，要安装keras模块

      .weights转.h5

      修改convert.py中

          

        需要避免[None，None，None，3]，设置为input_layer = Input(shape=(229, 229, 3))，tensorflow中为静态图，输入要确定。

        运行

  python convert.py yolov3-obj.cfg latest.weights latest.h5

        3.环境：TensorFlow2.0

import tensorflow as tf
converter = tf.lite.TFLiteConverter.from_keras_model_file(‘latest.h5‘)
tflite_model = converter.convert()
open("latest.tflite", "wb").write(tflite_model)

         生成后验证是否正确识别即可