訓練deepsort外觀度量模型cosine_metric_learning

語言: CN / TW / HK

訓練外觀度量模型cosine_metric_learning

本文描述訓練cosine metric,並匯出save model格式檔案,然後轉換成海思NNIE格式wk檔案的過程;

cosine_metric_learning 是多目標追蹤演算法deepsort裡用來生成影象特徵的CNN模型;

程式碼來源:https://github.com/nwojke/cosine_metric_learning.git

修改後的程式碼:https://gitee.com/duckship/cosine_metric_learning

原始碼是用TensorFlow1.x 的slim 寫的,現在安裝了TensorFlow2.0,不再內建slim,需要另外安裝第三方的庫,並對原始碼做些修改;

安裝tf-slim

pip install --upgrade tf_slim

匯入庫的時候需要做如下修改:

# import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

# import tensorflow.contrib.slim as slim
import tf_slim as slim

替換不支援的操作/OP

有些操作雖然caffe 和 NNIE 都支援,但是轉換工具ONNXToCaffe不支援,比如ELU, FC全連線層,全連線是因為ONNX不支援, 會用MUL和SUM替代,但是轉換工具ONNXToCaffe還不支援SUM操作;

  • ELU替換為relu
  • 將全連線替換為卷積
network = slim.conv2d(
                    network, feature_dim, kernel_size, stride=1, activation_fn=None,
                    padding="VALID", normalizer_fn=None, scope="fc1",
                    weights_initializer=conv_weight_init, biases_initializer=conv_bias_init,
                    weights_regularizer=conv_regularizer)

訓練方法,參照程式碼readme

python train_market1501.py --dataset_dir=./Market-1501-v15.09.15/ --loss_mode=cosine-softmax --log_dir=./train_data  --run_id=cosine-softmax

匯出checkpoint和saved_model格式資料

  • 修改train_app.py檔案的finalize函式程式碼,在函式末尾新增如下程式碼:
# save model
save_model_path = os.path.join(output_dir, 'save_model')
prediction_signature_dict = tf.saved_model.signature_def_utils.build_signature_def( 
    inputs={'input': tf.saved_model.utils.build_tensor_info(input_var)},
    outputs={'output': tf.saved_model.utils.build_tensor_info(output_tensor)},
    method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME
)
builder = tf.saved_model.builder.SavedModelBuilder(save_model_path)
builder.add_meta_graph_and_variables(session,
                                    ['serve'],
                                    {'signature': prediction_signature_dict})
builder.save()
  • 匯出命令
python .\train_market1501.py --mode=finalize --restore_path=<checkpoint path prefix>

restore_path 是訓練時產生的checkpoint,它不是一個完整的路徑,而是一個checkpoint檔案組的字首,比如: ./train_data/cosine-softmax/model.ckpt-556624

save model格式轉換為ONNX格式

  • 安裝onnx轉換工具
pip install -U tf2onnx
  • 轉換命令
python -m tf2onnx.convert  --saved-model  export/save_model --output export/cosine_metric_learning.onnx --opset 9

轉換命令的詳細使用方法請參與轉換工具的原始碼readme:https://github.com/onnx/tensorflow-onnx

ONNX格式轉換為caffe格式

  • 依賴條件:安裝好caffe的環境,支援所需要的層(操作、運算元);</br> caffe環境使用docker搭建;
  • 下載ONNXToCaffe 原始碼:[email protected]:duckship/onnxto-caffe.git
  • 簡化onnx模型,可以去掉一些caffe不支援的操作(如果報錯找不到檔案,請使用絕對路徑)
python3 -m onnxsim /workspace/docker-share/model/cosine_metric_learning.onnx /workspace/docker-share/model/cosine_metric_learning-sim.onnx
  • 轉換命令
python3 convertCaffe.py /workspace/docker-share/model/cosine_metric_learning-sim.onnx \
                        /workspace/docker-share/model/cosine_metric_learning.prototxt \
                        /workspace/docker-share/model/cosine_metric_learning.caffemodel

需要在ONNXToCaffe目錄執行,如果報錯找不到檔案,請使用絕對路徑;

caffe模型轉為NNIE指令檔案wk

修改上一步生成的protxt檔案

  • 修改輸入維度,使之符合NNIE要求:
layer {
 name: "input"
 type: "Input"
 top: "input"
 input_param {
   shape {
     dim: 1
     dim: 3
     dim: 128
     dim: 64
   }
 }
}
  • 去掉Permute

nnie_mapper cfg 檔案

檔名隨意,比如cosine_metric_learning.cfg

[prototxt_file] ./caffemodel/cosine_metric_learning.prototxt
[caffemodel_file] ./caffemodel/cosine_metric_learning.caffemodel
[batch_num] 1
[net_type] 0
[sparse_rate] 0
[compile_mode] 0
[is_simulation] 0
[log_level] 0
[instruction_name] ./nnie_wk/cosine_metric_learning
[RGB_order] BGR
[data_scale] 0.0039062
[internal_stride] 16
[image_list] ./images/image_ref_list.txt
[image_type] 3
[mean_file] null
[norm_type] 3

注意image_list欄位,可以在訓練圖片中隨機複製一些圖片到相應的檔案;

執行轉換命令

nnie_mapper_12 cosine_metric_learning.cfg