tensorflow 輸出pre_trained模型網絡層名稱以及數據

在模型訓練或者模型使用過程中，往往要獲取相關層的名稱或者特征，下面介紹如何讀取pre_trained model的各層數據以及各層的命名，以inception-v3模型為例。

• 1、預訓練模型下載

import numpy as np  
import tensorflow as tf  
import download  
from cache import cache  
import os  
import sys  

########################################################################  
# 壓縮包地址.  
data_url = "http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz"  

# 數據保存地址.  
data_dir = "inception/"  

# ImageNet 各個分類的名稱. (Downloaded)  
path_uid_to_cls = "imagenet_2012_challenge_label_map_proto.pbtxt"  

# File containing the mappings between uid and string. (Downloaded)  
path_uid_to_name = "imagenet_synset_to_human_label_map.txt"  

# 網絡層定義. (Downloaded)  
path_graph_def = "classify_image_graph_def.pb"  

########################################################################  


def maybe_download():  
    """ 
   如果inception-v3模型不存在就下載，大概85M. 
    """  

    print("Downloading Inception v3 Model ...")  
    download.maybe_download_and_extract(url=data_url, download_dir=data_dir)  

• 2、定義inception類

class Inception:
    """
    預訓練好的inception-v3包含1000種分類.
    """
    # 數據層.
    tensor_name_input_jpeg = "DecodeJpeg/contents:0"

    # resize后的數據.
    tensor_name_resized_image = "ResizeBilinear:0"
    # softmax層的名字.
    tensor_name_softmax_logits = "softmax/logits:0"

    # 最后一層的池化.
    tensor_name_transfer_layer = "pool_3:0"

    def __init__(self):

        # 創建tensorflow計算圖.
        self.graph = tf.Graph()

        # 將新的計算圖設置為默認圖.
        with self.graph.as_default():

            # 打開pre_trained模型.
            path = os.path.join(data_dir, path_graph_def)
            with tf.gfile.FastGFile(path, 'rb') as file:
                # 復制定義好的計算圖到新的圖中，先創建一個空的圖.
                graph_def = tf.GraphDef()

                # 加載proto-buf中的模型.
                graph_def.ParseFromString(file.read())

                # 最后復制pre-def圖的到默認圖中.
                tf.import_graph_def(graph_def, name='')

                # 完成從proto-buf的加載.

        # 獲取最后softmax層特征數據.
        self.y_logits = self.graph.get_tensor_by_name(self.tensor_name_softmax_logits)

        # 獲取計算圖最后一層的數據,可以更改對應名稱.
        self.transfer_layer = self.graph.get_tensor_by_name(self.tensor_name_transfer_layer)

        # 獲取最后一層的長度.
        self.transfer_len = self.transfer_layer.get_shape()[3]

        # 創建會話執行圖.
        self.session = tf.Session(graph=self.graph)

    def close(self):
        """
        關閉會話.
        """

        self.session.close()

    def _create_feed_dict(self, image_path=None, image=None):
        """
        """
        if image is not None:
            # Image is passed in as a 3-dim array that is already decoded.
            feed_dict = {self.tensor_name_input_image: image}

        elif image_path is not None:
            # Read the jpeg-image as an array of bytes.
            image_data = tf.gfile.FastGFile(image_path, 'rb').read()

            # Image is passed in as a jpeg-encoded image.
            feed_dict = {self.tensor_name_input_jpeg: image_data}

        else:
            raise ValueError("Either image or image_path must be set.")

        return feed_dict
    def transfer_values(self, image_path=None, image=None):
        """
        計算對應層數據

        :param image_path:
            輸入圖像路徑.

        :param image:
            輸入圖像數據.

        :return:
            對應層數據.
        """

        # Create a feed-dict for the TensorFlow graph with the input image.
        feed_dict = self._create_feed_dict(image_path=image_path, image=image)

        transfer_values = self.session.run(self.transfer_layer, feed_dict=feed_dict)

        # 變成一維數據輸出
        transfer_values = np.squeeze(transfer_values)

        return transfer_values

 

• 3、Demo

model = Inception()
#查看模型各層的名字
names = [op.name for op in model.graph.get_operations()]

values = model.transfer_values(image_path = "image_path")

print names
print values