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

• 主要思路 

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

 

• 步骤：

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

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

from models import *

model = Darknet("cfg/yolov3-obj.cfg")
#load_darknet_weights(model,"weights/latest.pt")
#save_weights(model,path='weights/latest.weights',cutoff=-1)

checkpoint = torch.load("weights/latest.pt", map_location='cpu')
model.load_state_dict(checkpoint['model'])
save_weights(model,path='weights/latest.weights',cutoff=-1)

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

      .weights转.h5

      修改输入层尺寸，input_layer = Input(shape=(416, 416, 3))，后续转换为tflite需要固定输入尺寸。不然会报错：

 

 

 

 

 

 

 

 

        运行

  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)

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