海思AI芯片(HI35xx):tensorflow轉caffemodel之模型參數轉換

摘要:
要把自己的模型進行移植,之前是後端的移植,最近前端也提了需求,前端一般都是用海思芯片(海思HI3516DV300),只支持caffe,所以爲了先測試時間得把tf的模型轉成caffemodel。這裏是將tf1.x轉爲caffemode,後續補全darknet轉爲caffemode

一、轉換ckpt轉caffemodel
轉換代碼:

#  coding=utf-8
#  Author      : AnnSun
#  Created date: 2020-06-23
#  

from __future__ import print_function, division
caffe_root = '/home/qif/smf/caffe/'   # 給爲你自己的caffe路徑
import sys
sys.path.insert(0, caffe_root + 'python')
import caffe
import numpy as np
import tensorflow as tf

deploy_proto = r"./yolov3_deploy_tf.prototxt"
caffe_model = r"./yolov3_helmet_tf.caffemodel"
# net = caffe.Net(deploy_proto, caffe_model, caffe.TEST)

# checkpoint爲模型保存地址
checkpoint_path = './models'
def tf_ckpt2caffemodel(deploy_proto, caffe_model, checkpoint_path):
    # 定義自己的net
    net = caffe.Net(deploy_proto, caffe.TEST)

    # 得到checkpoint文件中所有的參數(名字,形狀)元組
    for var_name, varshape in tf.contrib.framework.list_variables(checkpoint_path):
        # 得到上述參數的值
        var = tf.contrib.framework.load_variable(checkpoint_path, var_name)
        # var_name爲變量的name scope
        # var 是該name scope對應的值
        # conv layer
        ndim = var.ndim
        str_var_name = str(var_name)
        layerName_split = str_var_name.split("/")
        print(layerName_split)
        if str_var_name.endswith("weights"):
            layer_name = "_".join(layerName_split[1:-1])
            if ndim == 4:
                v_4d = np.transpose(var, [3, 2, 0, 1])
                net.params[layer_name][0].data[...] = v_4d
                print(var_name)
                print("conv weights {} succefully".format(layer_name))
                print(v_4d)

            else:
                v_2d = np.transpose(var, [1, 0])
                net.params[layer_name][0].data[...] = v_2d
                print("fc weights {} succefully".format(layer_name))
                print(v_2d)

        elif str_var_name.endswith('biases'):  # .strip("/biases:0")
            layer_name = "_".join(layerName_split[1:-1])
            net.params[layer_name][1].data[...] = var
            print("{} biases succefully".format(layer_name))
            print(var)


        elif layerName_split[-2] == "BatchNorm":
            # BN
            bn_name = "_".join(layerName_split[1:-1])
            if str_var_name.endswith('moving_mean'):
                net.params[bn_name][0].data[...] = var
                print("{} BN moving_mean succefully".format(bn_name))
                print(var)


            elif str_var_name.endswith('moving_variance'):
                net.params[bn_name][1].data[...] = var + 1e-3
                net.params[bn_name][2].data[...] = np.array([1.0])
                print("{} BN moving_variance succefully".format(bn_name))
                print(var)

            # scale
            elif str_var_name.endswith('beta'): # offset
                layer_name = "_".join(layerName_split[1:-2]) + "_scale"
                net.params[layer_name][1].data[...] = var
                print("{} Scale beta succefully".format(layer_name))
                print(var)


            elif str_var_name.endswith('gamma'):
                layer_name = "_".join(layerName_split[1:-2]) + "_scale"
                net.params[layer_name][0].data[...] = var
                print("{} Scale gamma succefully".format(layer_name))
                print(var)
    # 保存caffemodel
    net.save(caffe_model)
    print("\n -- caffeModel Finished. -- \n")


def show_TF_param(checkpoint_path):
    with open("tf1_yolov3_prama.txt", "w") as fw:
        # 得到checkpoint文件中所有的參數(名字,形狀)元組
        for var_name, varshape in tf.contrib.framework.list_variables(checkpoint_path):
            # 得到上述參數的值
            var = tf.contrib.framework.load_variable(checkpoint_path, var_name)
            # var_name爲變量的name scope
            # var 是該name scope對應的值
            info_str = str(var_name) + ":" + str(varshape) + "\n"
            fw.write(info_str)
            fw.write(str(var) + "\n" + "\n")
            print(var.ndim)
            print(var_name, " : ", varshape)
            print(var)

def show_caffe_param(deploy_proto, caffe_model):
    net = caffe.Net(deploy_proto, caffe_model, caffe.TEST)
    with open("caffe1_yolov3_prama.txt", "w") as fw:
        for layer_name, param in net.params.items():
            # print(layer_name + ": " + str(len(param)))
            for i in range(len(param)):
                info_str = layer_name + ': ' + str(i) + "/" + str(len(param)) + str(param[i].data.shape) + str(
                    param[i].data.ndim)
                print(info_str)
                fw.write(info_str + "\n")
                fw.write(str(param[i].data))
                fw.write("\n\n")

if __name__ == '__main__':
    # caffe
    deploy_proto = r"./yolov3_deploy_tf.prototxt"
    caffe_model = r"./yolov3_helmet_tf.caffemodel"
    # tf113(slim)
    # checkpoint爲模型保存地址
    checkpoint_path = './models'

    tfckpt_to_caffemodel = True
    isShowCaffe = True
    isShowTF = False

    if tfckpt_to_caffemodel:
        tf_ckpt2caffemodel(deploy_proto, caffe_model, checkpoint_path)
    elif isShowCaffe:
        show_caffe_param(deploy_proto, caffe_model)
    elif isShowTF:
        show_TF_param(checkpoint_path)

函數解析說明:
tf_ckpt2caffemodel():將tf的參數轉爲caffemodel的參數
show_caffe_param():轉換後的caffemodel的參數
show_TF_param():需要轉換的TF的模型中的參數

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章