2018年10月8日 星期一

把 Darknet 模型轉換成 CoreML 模型

要把 Darknet 模型轉換成 CoreML 模型,先用 Darkflow 把權重儲存成 TensorFlow PB 檔:
$ ./flow --model yolo-c3.cfg --load yolo-c3.weights --savepb
然後以 tfcoreml 做轉換。由於 tfcoreml 使用到不同的軟件版本組合,所以最好是用 Conda 之類的虛擬環境把軟件獨立出來:
$ git clone https://github.com/tf-coreml/tf-coreml.git
$ cd tf-coreml/
$ conda create --name tf-coreml python=3.6
$ source activate tf-coreml
$ pip install -e .
安裝好所需軟件版本後,下一步是正式轉換。把 darkflow/built_graph/yolo-c3.pb 拷到 tf-coreml 目錄,並進入 Python:
$ python
輸入以下 Python 程序。留意把下面「kerasModelPath」的值改為自己的 PB 檔路徑:
import tfcoreml as tf_converter
import tensorflow as tf

##----------------------------------------------------------------------------------------
##  We load the protobuf file from the disk and parse it to retrieve the unserialized graph_def
def load_graph(frozen_graph_filename):
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        
    # Then, we import the graph_def into a new Graph and return it 
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")
    return graph

##----------------------------------------------------------------------------------------
##  Load Keras model
kerasModelPath = 'yolo-c3.pb'
graph = load_graph(kerasModelPath)
for op in graph.get_operations(): 
    print (op.name)

##----------------------------------------------------------------------------------------
##  Convert Keras model to Core ML model
##  output_feature_names: the output node name we get from the previouse step
##  image_input_names: CoreML allows image as the input, the only thing we need to do is to set which node is the image input node 
##  input_name_shape_dict: the input node name we get from the previous step, and check the cfg file to know the exact input shape size
##  is_bgr: the channel order is by BGR instead of RGB
##  image_scale: the weights is already normalized in the range from 0 to 1
coreml_model = tf_converter.convert(tf_model_path=kerasModelPath, mlmodel_path='yolo.mlmodel', output_feature_names=['grid'], image_input_names= ['image'], input_name_shape_dict={'image': [1, 416, 416, 3]}, is_bgr=True, image_scale=1/255.0)
完成後便會得到 tf-coreml/yolo-c3.mlmodel 模型檔。