YOLOv4部署－keras权重转tflite (h5 weights -＞ pb -＞ tflite)

一、相关环境介绍

1、网络训练的环境

• tensorflow-gpu==1.4.0
• Keras==2.1.5
• python==3.6.2

２、h5->pb->tflite的环境

• tensorflow==2.2.0
• python==3.6.2

二、转换过程

１、加载h5权重，并转换为pb

• 修改yolov4网络结构中keras的导包命令，tensorflow=2.2.0中自带的有keras

  # 将所有的
  from keras import *
  from keras.* import *
  # 修改过为：
  from tensorflow.keras import *
  from tensorflow.keras.* import *
  

• 加载权重和网络结构

  from tensorflow.keras.models import Model
  import tensorflow as tf
  class YOLO(object):
  _defaults = {
  ...
  }
  @classmethod
  def get_defaults(cls, n):
  if n in cls._defaults:
  return cls._defaults[n]
  else:
  return "Unrecognized attribute name '" + n + "'"
  #---------------------------------------------------#
  #
  初始化yolo
  #---------------------------------------------------#
  def __init__(self, **kwargs):
  self.__dict__.update(self._defaults)
  for name, value in kwargs.items():
  setattr(self, name, value)
  self._defaults[name] = value
  self.class_names, self.num_classes = get_classes(self.classes_path)
  self.anchors, self.num_anchors
  = get_anchors(self.anchors_path)
  self.generate()
  #---------------------------------------------------#
  #
  载入模型
  #---------------------------------------------------#
  def generate(self):
  model_path = os.path.expanduser(self.model_path)
  assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'
  self.model = yolo_body([640, 640, 3], ...)
  self.model.load_weights(self.model_path)
  outputs = Lambda(
  DecodeBox,
  output_shape = (1,),
  name = 'yolo_eval',
  arguments = {...}
  )(self.model.output)
  # self.model.output输出是三个特征层
  self.yolo_model = Model(self.model.input, outputs)
  tf.saved_model.save(self.yolo_model, "yolo_tflite/yolov4")
  

２、pb文件转换为tflite

 converter = tf.lite.TFLiteConverter.from_saved_model('yolo_tflite/yolov4')
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
converter.allow_custom_ops = True
tflite_model = converter.convert()
open('yolo_tflite/yolo_fp32.tflite', 'wb').write(tflite_model)

最后

以上就是知性胡萝卜为你收集整理的YOLOv4部署－keras权重转tflite (h5 weights -＞ pb -＞ tflite)的全部内容，希望文章能够帮你解决YOLOv4部署－keras权重转tflite (h5 weights -＞ pb -＞ tflite)所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。