Keras训练的h5模型转Tensorflow pb模型OpenCV可调用

环境：Keras 2.2.4

           Tensorflow-gpu 1.12

这里的h5模型是由keras训练保存的，注意不是tf.keras。因为如果是tf.keras训练生成的模型，那可能没这么多坑了。

也就是用的

from keras.layers import Dense, GlobalAveragePooling2D, Dropout, Input, Flatten, Conv2D,Softmax
   
   

这样构建的网络。

注意保存模型时要使用model.save()，如果使用model.save_weights()仅保存了权重，要恢复还要知道网络结构。

这里是保存的h5模型，注意到有BatchNormalization这样的操作，多半的坑都是因为它。

坑：

加载h5模型时使用的是

   
   

     
     

      
      

     
     
     
     

      
      

       
       from keras.models 
       
       import load_model
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       my_model = load_model(
       
       'my_model.h5')
      
      
     
     
   
   

然后按照后面的流程也可以成功得到pb模型文件，但是opencv不能成功加载模型，看一下这个模型

和BatchNormalization有关的节点明显改变了。

看一下两个模型在BatchNormalization这个节点的不同。

如果用opencv的dnn来加载这个pb文件，会报错

cv2.error: OpenCV(4.1.1) C:\projects\opencv-python\opencv\modules\dnn\src\tensorflow\tf_importer.cpp:582: error: (-2:Unspecified error) Input [batch_normalization_1/ones_like] for node [batch_normalization_1/FusedBatchNorm_1] not found in function 'cv::dnn::dnn4_v20190621::`anonymous-namespace'::TFImporter::getConstBlob'
   
   

从这个log也可以推断是和BatchNormalization有关。

排坑：

 一定要用tensorflow.keras加载h5文件。

   
   

     
     

      
      

     
     
     
     

      
      

       
       from tensorflow.keras 
       
       as keras
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       my_model = keras.models.load_model(
       
       'my_model.h5')
      
      
     
     
   
   

因为如果用原生keras加载模型会使用keras.engine.sequential.Sequential，这可能会导致一些object无法直接转换成tensorflow pb文件。

可以看到用tensorflow keras保存的模型格式已经改变了。

所以完整的流程应该是:

1. keras训练，保存h5文件
2. tensorflow keras加载h5文件
3. 转换成tensorflow pb文件
4. opencv dnn调用pb文件

   
   

     
     

      
      

     
     
     
     

      
      

       
       import tensorflow.keras 
       
       as keras
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       import tensorflow 
       
       as tf
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       import os
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       #这个函数参考自网上
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
      
      
     
     

     
     

      
      

     
     
     
     

      
      
    
       
       """
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           Freezes the state of a session into a pruned computation graph.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           Creates a new computation graph where variable nodes are replaced by
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           constants taking their current value in the session. The new graph will be
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           pruned so subgraphs that are not necessary to compute the requested
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           outputs are removed.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           @param session The TensorFlow session to be frozen.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           @param keep_var_names A list of variable names that should not be frozen,
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                                 or None to freeze all the variables in the graph.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           @param output_names Names of the relevant graph outputs.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           @param clear_devices Remove the device directives from the graph for better portability.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           @return The frozen graph definition.
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           """
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           graph = session.graph
      
      
     
     

     
     

      
      

     
     
     
     

      
      
    
       
       with graph.as_default():
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
               freeze_var_names = list(set(v.op.name 
       
       for v 
       
       in tf.global_variables()).difference(keep_var_names 
       
       or []))
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
               output_names = output_names 
       
       or []
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
               output_names += [v.op.name 
       
       for v 
       
       in tf.global_variables()]
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
               input_graph_def = graph.as_graph_def()
      
      
     
     

     
     

      
      

     
     
     
     

      
      
        
       
       if clear_devices:
      
      
     
     

     
     

      
      

     
     
     
     

      
      
            
       
       for node 
       
       in input_graph_def.node:
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                       node.device = 
       
       ''
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
               frozen_graph = tf.graph_util.convert_variables_to_constants(
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
                   session, input_graph_def, output_names, freeze_var_names)
      
      
     
     

     
     

      
      

     
     
     
     

      
      
        
       
       return frozen_graph
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
       if __name__ == 
       
       '__main__':
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           input_path = 
       
       'D:\\training'
      
      
     
     

     
     

      
      

     
     
     
     

      
      
    
       
       #keras训练保存的h5文件
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           input_file = 
       
       'my_model.h5'
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           weight_file_path = os.path.join(input_path, input_file)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           output_graph_name = weight_file[:
       
       -3] + 
       
       '.pb'
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      
    
       
       #  加载模型
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           keras.backend.set_learning_phase(
       
       0)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           h5_model = keras.models.load_model(weight_file_path)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           frozen_graph = freeze_session(keras.backend.get_session(), output_names=[out.op.name 
       
       for out 
       
       in h5_model.outputs])
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           tf.train.write_graph(frozen_graph, input_path, output_graph_name, as_text=
       
       False)
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           print(
       
       'Finished')
      
      
     
     

     
     

      
      

     
     
     
     

      
      
 
      
      
     
     

     
     

      
      

     
     
     
     

      
      
    
       
       import cv2
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           model = cv2.dnn.readNetFromTensorflow(
       
       "D:\\training\\my_model.pb")
      
      
     
     

     
     

      
      

     
     
     
     

      
      

       
           print(
       
       'Load')
      
      
     
     
   
   

通过这种方式就可以正确转换成pb文件了，并且这个pb文件opencv也是可以成功调用的了。

但是转换之后模型变成了

似乎是有一些在推理中不必要的节点。

可以通过下面的方式对pb模型进行进一步优化

python -m tensorflow.python.tools.optimize_for_inference --input my_model.pb --output my_model_opt.pb --input_names=input_1 --output_names=dense_1/Softmax
   
   

优化之后模型的体积小了不少，那些多出来的节点也没有了，另外BatchNormalization这样在推理中不需要的节点也没有了。