獲取和生成基於TensorFlow的MobilNet預訓練模型

bashan16045發表於2020-11-03
模型

從TensorFlow-Slim image classification model library中找到我們要使用的MobileNet_v2_1.0_224^*（1.0-阿爾法值，244-輸入圖片shape），下載mobilenet_v2_1.0_224.tgz檔案，其中內含：

  • mobilenet_v2_1.0_224.ckpt.data-00000-of-00001
  • mobilenet_v2_1.0_224.ckpt.index
  • mobilenet_v2_1.0_224.ckpt.meta
  • mobilenet_v2_1.0_224.tflite
  • mobilenet_v2_1.0_224_eval.pbtxt
  • mobilenet_v2_1.0_224_frozen.pb
  • mobilenet_v2_1.0_224_info.txt

mobilenet_v2_1.0_224.ckpt.data-00000-of-00001儲存模型權重，mobilenet_v2_1.0_224.ckpt.meta儲存模型圖的流程，mobilenet_v2_1.0_224.ckpt.index儲存模型結構的變數與引數間的索引對應關係，mobilenet_v2_1.0_224_frozen.pb凍結的PB檔案，mobilenet_v2_1.0_224.tflite凍結的tflite檔案。

在Test6_mobilenet下新建目錄pretain_model，將以下三個檔案放入其中

  • mobilenet_v2_1.0_224.ckpt.data-00000-of-00001
  • mobilenet_v2_1.0_224.ckpt.meta
  • mobilenet_v2_1.0_224.ckpt.index

生成預訓練模型程式碼，執行後在當前路徑下生成兩個檔案：

  • pretrain_weights.ckpt.data-00000-of-00001
  • pretrain_weights.ckpt.index
import tensorflow as tf


def rename_var(ckpt_path, new_ckpt_path, num_classes=5):
    with tf.Graph().as_default(), tf.compat.v1.Session().as_default() as sess:
        var_list = tf.train.list_variables(ckpt_path)
        new_var_list = []

        for var_name, shape in var_list:
            # print(var_name)
            #filter 不需要的層結構
            if var_name in except_list:
                continue
            #filter所有有關優化器info
            if "RMSProp" in var_name or "Exponential" in var_name:
                continue
            #層結構的名稱轉換
            var = tf.train.load_variable(ckpt_path, var_name)
            new_var_name = var_name.replace('MobilenetV2/', "")
            new_var_name = new_var_name.replace("/expand/weights", "/expand/Conv2d/weights")
            new_var_name = new_var_name.replace("Conv/weights", "Conv/Conv2d/kernel")
            new_var_name = new_var_name.replace("Conv_1/weights", "Conv_1/Conv2d/kernel")
            new_var_name = new_var_name.replace("weights", "kernel")
            new_var_name = new_var_name.replace("biases", "bias")

            first_word = new_var_name.split('/')[0]
            if "expanded_conv" in first_word:
                last_word = first_word.split('expanded_conv')[-1]
                if len(last_word) > 0:
                    new_word = "inverted_residual" + last_word + "/expanded_conv/"
                else:
                    new_word = "inverted_residual/expanded_conv/"
                new_var_name = new_word + new_var_name.split('/', maxsplit=1)[-1]
            print(new_var_name)
            re_var = tf.Variable(var, name=new_var_name)
            new_var_list.append(re_var)

        re_var = tf.Variable(tf.keras.initializers.he_uniform()([1280, num_classes]), name="Logits/kernel")
        new_var_list.append(re_var)
        re_var = tf.Variable(tf.keras.initializers.he_uniform()([num_classes]), name="Logits/bias")

        new_var_list.append(re_var)
        tf.keras.initializers.he_uniform()
        saver = tf.compat.v1.train.Saver(new_var_list)
        sess.run(tf.compat.v1.global_variables_initializer())
        saver.save(sess, save_path=new_ckpt_path, write_meta_graph=False, write_state=False)

# 不需要的層結構
# 'MobilenetV2/Logits/Conv2d_1c_1x1/biases', 'MobilenetV2/Logits/Conv2d_1c_1x1/weights'
# MobilenetV2的全連線層偏置和權重
except_list = ['global_step', 'MobilenetV2/Logits/Conv2d_1c_1x1/biases', 'MobilenetV2/Logits/Conv2d_1c_1x1/weights']
ckpt_path = './pretain_model/mobilenet_v2_1.0_224.ckpt'
new_ckpt_path = './pretrain_weights.ckpt'
num_classes = 5
rename_var(ckpt_path, new_ckpt_path, num_classes)

 

相關文章