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)

         生成后驗證是否正確識別即可