tensorflow 输出pre_trained模型网络层名称以及数据

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