TensorFlow1.x如何使用官方预训练模型

谁借莪１个温暖的怀抱￠ 2022-09-08 14:52 203阅读 0赞

### 文章目录 ###

*  模型下载
 *  辅助函数
 *  模型使用
 *   *  1. 使用官方提供的模型文件加载模型
     *  2. 魔改官方的模型文件
     *  3. 使用NewCheckpointReader并自定义模型

环境：tensorflow1.13  
模型：使用vgg19作为例子

注意：本文档的结果是使用CPU模式跑出来的，因为显卡是30系，系统是win10，没法在tf1.13版本下使用GPU模式。所以如果读者使用GPU模式跑出来的结果与本文档有些许出入，应该是正常现象。

文档背景是为了在tensorflow1.x中使用预训练的vgg模型计算vgg loss。

# 模型下载 #

官方的预训练模型在tensorflow的model仓库中，完整路径是`tensorflow/models/research/slim`，请注意选择tf1.13分支：  
[https://github.com/tensorflow/models/tree/r1.13.0/research/slim][https_github.com_tensorflow_models_tree_r1.13.0_research_slim]

vgg19的模型下载地址：  
[http://download.tensorflow.org/models/vgg\_19\_2016\_08\_28.tar.gz][http_download.tensorflow.org_models_vgg_19_2016_08_28.tar.gz]  
解压后文件名是`vgg_19.ckpt`。

# 辅助函数 #

在下面的代码中，为了对比不同的模型使用方式是否带来一致的结果，写了一个可视化feature map的辅助函数来方便对比。下面的代码将直接使用，而不再赘述，读者如果想要尝试本文档代码，请自行粘贴过去。

def visualize_feature_map(feature_map,
                              col_nums=None,
                              gap_value=0.5,
                              gap_width=10,
                              gap_height=10):
        """ Visualize feature map in one image. Parameters ---------- feature_map: numpy array, shape is (height, width, channel) col_nums: number of feature map columns gap_value: value for feature map gap gap_width: width of gap gap_height: height of gap Returns ------- image: image to show feature map """
        eps = 1e-6
    
        if feature_map.ndim == 4:
            if feature_map.shape[0] == 1:
                feature_map = np.squeeze(feature_map)
            else:
                raise ValueError("feature map must be 3 dims ndarray (height, "
                                 "width, channel) or 4 dims ndarray whose shape "
                                 "must be (1, height, width, channel)")
    
        # compute col_nums (if not set) and row_nums
        height, width, channel = feature_map.shape
        if col_nums is None:
            col_nums = int(round(np.sqrt(channel)))
        row_nums = int(np.ceil(channel / col_nums))
    
        # compute final image width and height
        image_width = col_nums * (width + gap_width) - gap_width
        image_height = row_nums * (height + gap_height) - gap_height
    
        image = np.ones(shape=(image_height, image_width),
                        dtype=feature_map.dtype) * gap_value
        cnt = 0
        while cnt < channel:
            row = cnt // col_nums
            col = cnt % col_nums
    
            row_beg = row * (height + gap_height)
            row_end = row_beg + height
            col_beg = col * (width + gap_width)
            col_end = col_beg + width
    
            image[row_beg:row_end, col_beg:col_end] = \
                feature_map[:, :, cnt] / (np.std(feature_map[:, :, cnt]) + eps)
            cnt += 1
    
        return image

# 模型使用 #

有三种典型的方式使用官方模型：

1.  使用官方提供的模型文件加载模型。  
    需要首先使用placeholder在计算图中定义模型，然后使用tf.train.Saver()的restore方法加载模型参数。这种方式要求我们新定义的模型的节点名和参数名必须与`vgg_19.ckpt`中保存的一致，这也是为什么建议直接使用官方的模型定义文件。
2.  魔改官方的模型文件。  
    在卷积部分，官方的模型文件只提供了relu层的feature map，有时候我们可能需要conv层的feature map，这时候就需要魔改。魔改后的节点名和参数名仍然必须与`vgg_19.ckpt`中保存的一致。
3.  使用NewCheckpointReader并自定义模型  
    使用`pywrap_tensorflow.NewCheckpointReader(model_path)`可以读取权重参数，然后自己重新定义模型结构并把权重参数赋值过去。采用这种办法可以**灵活**定义模型结构和节点名称，但是代码写起来比较麻烦。（官方定义的模型文件只能拿到relu层的feature map，拿不到conv层的，所以灵活性欠佳）

## 1. 使用官方提供的模型文件加载模型 ##

官方的模型定义文件路径（网址）：  
https://github.com/tensorflow/models/blob/r1.13.0/research/slim/nets/vgg.py  
找到vgg19的定义如下：

def vgg_19(inputs,
               num_classes=1000,
               is_training=True,
               dropout_keep_prob=0.5,
               spatial_squeeze=True,
               scope='vgg_19',
               fc_conv_padding='VALID',
               global_pool=False):
        """ Oxford Net VGG 19-Layers version E Example. Note: All the fully_connected layers have been transformed to conv2d layers. To use in classification mode, resize input to 224x224. Args: inputs: a tensor of size [batch_size, height, width, channels]. num_classes: number of predicted classes. If 0 or None, the logits layer is omitted and the input features to the logits layer are returned instead. is_training: whether or not the model is being trained. dropout_keep_prob: the probability that activations are kept in the dropout layers during training. spatial_squeeze: whether or not should squeeze the spatial dimensions of the outputs. Useful to remove unnecessary dimensions for classification. scope: Optional scope for the variables. fc_conv_padding: the type of padding to use for the fully connected layer that is implemented as a convolutional layer. Use 'SAME' padding if you are applying the network in a fully convolutional manner and want to get a prediction map downsampled by a factor of 32 as an output. Otherwise, the output prediction map will be (input / 32) - 6 in case of 'VALID' padding. global_pool: Optional boolean flag. If True, the input to the classification layer is avgpooled to size 1x1, for any input size. (This is not part of the original VGG architecture.) Returns: net: the output of the logits layer (if num_classes is a non-zero integer), or the non-dropped-out input to the logits layer (if num_classes is 0 or None). end_points: a dict of tensors with intermediate activations. """
        with tf.variable_scope(scope, 'vgg_19', [inputs]) as sc:
            end_points_collection = sc.original_name_scope + '_end_points'
            # Collect outputs for conv2d, fully_connected and max_pool2d.
            with slim.arg_scope(
                    [slim.conv2d, slim.fully_connected, slim.max_pool2d],
                    outputs_collections=end_points_collection):
                net = slim.repeat(inputs, 2, slim.conv2d, 64, [3, 3],
                                  scope='conv1')
                net = slim.max_pool2d(net, [2, 2], scope='pool1')
                net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3], scope='conv2')
                net = slim.max_pool2d(net, [2, 2], scope='pool2')
                net = slim.repeat(net, 4, slim.conv2d, 256, [3, 3], scope='conv3')
                net = slim.max_pool2d(net, [2, 2], scope='pool3')
                net = slim.repeat(net, 4, slim.conv2d, 512, [3, 3], scope='conv4')
                net = slim.max_pool2d(net, [2, 2], scope='pool4')
                net = slim.repeat(net, 4, slim.conv2d, 512, [3, 3], scope='conv5')
                net = slim.max_pool2d(net, [2, 2], scope='pool5')
    
                # Use conv2d instead of fully_connected layers.
                net = slim.conv2d(net, 4096, [7, 7], padding=fc_conv_padding,
                                  scope='fc6')
                net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
                                   scope='dropout6')
                net = slim.conv2d(net, 4096, [1, 1], scope='fc7')
                # Convert end_points_collection into a end_point dict.
                end_points = slim.utils.convert_collection_to_dict(
                    end_points_collection)
                if global_pool:
                    net = tf.reduce_mean(net, [1, 2], keep_dims=True,
                                         name='global_pool')
                    end_points['global_pool'] = net
                if num_classes:
                    net = slim.dropout(net, dropout_keep_prob,
                                       is_training=is_training,
                                       scope='dropout7')
                    net = slim.conv2d(net, num_classes, [1, 1],
                                      activation_fn=None,
                                      normalizer_fn=None,
                                      scope='fc8')
                    if spatial_squeeze:
                        net = tf.squeeze(net, [1, 2], name='fc8/squeezed')
                    end_points[sc.name + '/fc8'] = net
                return net, end_points

关于上述模型定义，有两个函数需要稍加解释：

*  slim.conv2d  
    `slim.conv2d`的定义在Pycharm中通过ctrl加单击的方式追进去后发现不对，真正的定义在下面路径的文件中：  
    D:\\Program\\anaconda3\\envs\\tf13\\Lib\\site-packages\\tensorflow\\contrib\\layers\\python\\layers\\layers.py  
    （请注意D:\\Program\\anaconda3\\envs\\tf13是我电脑上tf1.13的环境路径，需要根据自己环境修改）
    
    其中第1117行的`def convolution2d`是函数定义及实现，3327行的`conv2d = convolution2d`起了一个短名字。`convolution2d`的参数列表中有`activation_fn=nn.relu`，所以这个卷积默认带有relu作为激活函数。
 *  slim.repeat  
    作用是把某个算子重复n次。函数实现跟slim.conv2d同一个文件，文件中的函数定义和部分解释如下：
    
        def repeat(inputs, repetitions, layer, *args, **kwargs):
          """Applies the same layer with the same arguments repeatedly. y = repeat(x, 3, conv2d, 64, [3, 3], scope='conv1') # It is equivalent to: x = conv2d(x, 64, [3, 3], scope='conv1/conv1_1') x = conv2d(x, 64, [3, 3], scope='conv1/conv1_2') y = conv2d(x, 64, [3, 3], scope='conv1/conv1_3') ...... """

使用脚本如下，注意其中`vgg_19`和`visualize_feature_map`两个函数在本文档上面已经出现过，并且比较长，所以下面的脚本中将其省略：

# -*- coding: utf-8 -*-
    import os
    import cv2
    import tensorflow as tf
    import numpy as np
    
    os.environ['CUDA_VISIBLE_DEVICES'] = "/gpu:0"
    slim = tf.contrib.slim
    
    
    def vgg_19(inputs,
               num_classes=1000,
               is_training=True,
               dropout_keep_prob=0.5,
               spatial_squeeze=True,
               scope='vgg_19',
               fc_conv_padding='VALID',
               global_pool=False):
        # 见本文档前面
        pass
    
    
    def visualize_feature_map(feature_map,
                              col_nums=None,
                              gap_value=0.5,
                              gap_width=10,
                              gap_height=10):
        # 见本文档前面
        pass
    
    
    def main():
        image_file = r'E:\images\lena512color.tiff'
        model_path = r'E:\pretrained_model\tf1x\vgg_19.ckpt'
        inputs_ = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
        outputs, feature_map_dict = vgg_19(inputs_,
                                           num_classes=0,
                                           is_training=False,
                                           global_pool=True)
    
        # print trainable variables
        for var in tf.trainable_variables():
            print(var)
    
        # load pretrained model
        saver = tf.train.Saver()
        sess = tf.Session()
        saver.restore(sess, model_path)
    
        # running test
        inputs = cv2.imread(image_file)
        inputs = np.expand_dims(inputs, axis=0)
        out, feature_maps = sess.run([outputs, feature_map_dict],
                                     feed_dict={ 
                                         inputs_: inputs,
                                     })
    
        # print shape of feature maps
        for key in feature_maps.keys():
            print(key, feature_maps.get(key).shape)
    
        feature_map = feature_maps.get('vgg_19/conv3/conv3_4')
        feature_map = np.squeeze(feature_map)
        image = visualize_feature_map(feature_map)
        image = np.clip(image * 255, 0, 255).astype(np.uint8)
        # cv2.imwrite('lena_feature_map_vgg_conv3_4.png', image)
    
        # print statistics for feature map
        for i in range(5):
            mean_val = np.mean(feature_map[:, :, i])
            std = np.std(feature_map[:, :, i])
            min_val = np.min(feature_map[:, :, i])
            max_val = np.max(feature_map[:, :, i])
            print(i + 1, " min=%.4f, max=%.4f, mean=%.4f, std=%.4f" % (
                min_val, max_val, mean_val, std))
    
        # print part of final global feature vector
        feature_vec = feature_maps.get('global_pool')
        feature_vec = np.squeeze(feature_vec)
        for i in range(10):
            print(feature_vec[i])
    
    
    if __name__ == '__main__':
        main()

对上述脚本需要做如下说明：

1.  vgg\_19的参数设置需要多加注意。  
    如果只是做推理的话，is\_training一定要设置成False；  
    我使用的目的是为了计算vgg loss，所以不需要全连接的部分，因此为了使全连接部分不要报错，把num\_classes设置成0，并且global\_pool设置为True。
2.  使用placeholder创建计算图以后才能restore权重。  
    脚本在计算图和restore权重的中间部分(`# print trainable variables`部分)打印了权重变量，包括名称，shape，dtype。
    
        <tf.Variable 'vgg_19/conv1/conv1_1/weights:0' shape=(3, 3, 3, 64) dtype=float32_ref>
        <tf.Variable 'vgg_19/conv1/conv1_1/biases:0' shape=(64,) dtype=float32_ref>
        <tf.Variable 'vgg_19/conv1/conv1_2/weights:0' shape=(3, 3, 64, 64) dtype=float32_ref>
        <tf.Variable 'vgg_19/conv1/conv1_2/biases:0' shape=(64,) dtype=float32_ref>
        <tf.Variable 'vgg_19/conv2/conv2_1/weights:0' shape=(3, 3, 64, 128) dtype=float32_ref>
        <tf.Variable 'vgg_19/conv2/conv2_1/biases:0' shape=(128,) dtype=float32_ref>
        ......
3.  `vgg_19`的输出有两个。  
    第一个很容易理解，就是网络推理的输出，但对于计算vgg loss而言没有用。  
    第二个输出以dict的形式保存了网络的feature map，dict的key是feature map的节点名，value是feature map的数值，这是计算vgg loss真正需要的东西。在`# print shape of feature maps`部分打印了feature map的name和shape，另外还把`conv3_4`画在一张图上，用于一些简单直观的测试和检查。
    
        vgg_19/conv1/conv1_1 (1, 512, 512, 64)
        vgg_19/conv1/conv1_2 (1, 512, 512, 64)
        vgg_19/pool1 (1, 256, 256, 64)
        vgg_19/conv2/conv2_1 (1, 256, 256, 128)
        vgg_19/conv2/conv2_2 (1, 256, 256, 128)
        vgg_19/pool2 (1, 128, 128, 128)
        ......
4.  打印feature map的一些统计数值，可以检查并确认以下事实 ：  
    feature map只有relu，没有conv，因为feature map的最小值都是0.0；  
    vgg出现在BN之前，所以网络中没有BN，导致feature map的数值很大（如果有BN的话数值一般不会超过5），因此计算vgg loss时根据具体情况一般需要乘上一个很小的权重系数。
    
        1  min=0.0000, max=9201.5811, mean=386.3745, std=737.3252
        2  min=0.0000, max=7389.5913, mean=1412.0540, std=616.6437
        3  min=0.0000, max=3323.7239, mean=400.2662, std=522.4063
        4  min=0.0000, max=4319.3765, mean=369.9904, std=644.4222
        5  min=0.0000, max=8997.2305, mean=905.1512, std=1288.8953
        ......
5.  打印一部分最后的feature vector，用于魔改模型定义函数后检查正确性：
    
        0.00055606366
        0.0
        0.0
        0.15579844
        0.0
        1.0548652
        0.0
        0.0
        0.05207316
        0.29752082
        ......

## 2. 魔改官方的模型文件 ##

魔改后的模型和测试代码如下，同样`visualize_feature_map`需要从上面粘贴过来：

# -*- coding: utf-8 -*-
    import os
    import cv2
    import tensorflow as tf
    import numpy as np
    
    slim = tf.contrib.slim
    
    
    def vgg19(inputs,
              num_classes=1000,
              is_training=True,
              dropout_keep_prob=0.5,
              spatial_squeeze=True,
              scope='vgg_19',
              fc_conv_padding='VALID',
              global_pool=False):
        with tf.variable_scope(scope, 'vgg_19', [inputs]) as sc:
            end_points_collection = sc.original_name_scope + '_end_points'
            # Collect outputs for conv2d, fully_connected and max_pool2d.
            with slim.arg_scope(
                    [slim.conv2d, slim.fully_connected, slim.max_pool2d],
                    outputs_collections=end_points_collection):
                # conv blocks are modified as follows
                net_config = [
                    [64, 2],
                    [128, 2],
                    [256, 4],
                    [512, 4],
                    [512, 4],
                ]  # [filters, blocks]
    
                net = inputs
                relu_dict = { }
                for i, config in enumerate(net_config):
                    filters = config[0]
                    for j in range(config[1]):
                        conv_scope = 'conv%d/conv%d_%d' % (i + 1, i + 1, j + 1)
                        relu_name = 'conv%d/relu%d_%d' % (i + 1, i + 1, j + 1)
                        net = slim.conv2d(net, filters, [3, 3],
                                          activation_fn=None,
                                          scope=conv_scope)
                        net = tf.nn.relu(net, name=relu_name)
                        relu_dict[net.op.name] = net
                    net = slim.max_pool2d(net, [2, 2], scope='pool%d' % (i + 1))
    
                # Use conv2d instead of fully_connected layers.
                net = slim.conv2d(net, 4096, [7, 7], padding=fc_conv_padding,
                                  scope='fc6')
                net = slim.dropout(net, dropout_keep_prob, is_training=is_training,
                                   scope='dropout6')
                net = slim.conv2d(net, 4096, [1, 1], scope='fc7')
                # Convert end_points_collection into a end_point dict.
                end_points = slim.utils.convert_collection_to_dict(
                    end_points_collection)
                if global_pool:
                    net = tf.reduce_mean(net, [1, 2], keep_dims=True,
                                         name='global_pool')
                    end_points['global_pool'] = net
                if num_classes:
                    net = slim.dropout(net, dropout_keep_prob,
                                       is_training=is_training,
                                       scope='dropout7')
                    net = slim.conv2d(net, num_classes, [1, 1],
                                      activation_fn=None,
                                      normalizer_fn=None,
                                      scope='fc8')
                    if spatial_squeeze:
                        net = tf.squeeze(net, [1, 2], name='fc8/squeezed')
                    end_points[sc.name + '/fc8'] = net
    
                end_points.update(relu_dict)
                return net, end_points
    
    
    def visualize_feature_map(feature_map,
                              col_nums=None,
                              gap_value=0.5,
                              gap_width=10,
                              gap_height=10):
        # 见本文档前面
        pass
    
    
    def main():
        image_file = r'D:\data\test_images\lena512color.tiff'
        model_path = r'E:\pretrained_model\tensorflow1.13\vgg_19.ckpt'
        inputs_ = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
        outputs, feature_map_dict = vgg19(inputs_,
                                          num_classes=0,
                                          is_training=False,
                                          global_pool=True)
    
        # check trainable variables
        for var in tf.trainable_variables():
            print(var)
    
        # load pretrained model
        saver = tf.train.Saver()
        sess = tf.Session()
        saver.restore(sess, model_path)
    
        # running test
        inputs = cv2.imread(image_file)
        inputs = np.expand_dims(inputs, axis=0)
        out, feature_maps = sess.run([outputs, feature_map_dict],
                                     feed_dict={ 
                                         inputs_: inputs,
                                     })
    
        # print shape of feature maps
        for key in feature_maps.keys():
            print(key, feature_maps.get(key).shape)
    
        feature_map = feature_maps.get('vgg_19/conv3/relu3_4')
        feature_map = np.squeeze(feature_map)
        image = visualize_feature_map(feature_map)
        image = np.clip(image * 255, 0, 255).astype(np.uint8)
        cv2.imwrite('lena_feature_map_vgg_relu3_4--2.png', image)
    
        # print statistics for relu3_4
        for i in range(5):
            mean_val = np.mean(feature_map[:, :, i])
            std = np.std(feature_map[:, :, i])
            min_val = np.min(feature_map[:, :, i])
            max_val = np.max(feature_map[:, :, i])
            print(i + 1, " min=%.4f, max=%.4f, mean=%.4f, std=%.4f" % (
                min_val, max_val, mean_val, std))
    
        # print statistics for conv3_4
        print('\n')
        feature_map = feature_maps.get('vgg_19/conv3/conv3_4')
        feature_map = np.squeeze(feature_map)
        for i in range(5):
            mean_val = np.mean(feature_map[:, :, i])
            std = np.std(feature_map[:, :, i])
            min_val = np.min(feature_map[:, :, i])
            max_val = np.max(feature_map[:, :, i])
            print(i + 1, " min=%.4f, max=%.4f, mean=%.4f, std=%.4f" % (
                min_val, max_val, mean_val, std))
    
        feature_vec = feature_maps.get('global_pool')
        feature_vec = np.squeeze(feature_vec)
        for i in range(10):
            print(feature_vec[i])
    
    
    if __name__ == '__main__':
        main()

说明如下：

1.  改模型结构的目的：将conv和relu分离开来，方便使用conv层的结果来作为vgg loss的输入。
2.  改模型定义函数的关键点：确保结构不能变；确保节点名不能变；由于原始代码不能直接收集relu层的feature map，所以需要自行定义一个dict进行收集。
3.  conv3\_4的feature map统计信息如下，可以看到min value已经出现负数，所以确实实现了与relu层的分离。
    
        1  min=-2909.7209, max=9201.5811, mean=92.2542, std=991.5225
        2  min=-431.0446, max=7389.5913, mean=1411.6982, std=617.5237
        3  min=-1092.2075, max=3323.7239, mean=339.5828, std=582.6731
        4  min=-2396.4478, max=4319.3765, mean=106.6278, std=852.0536
        5  min=-3547.4551, max=8997.2305, mean=699.2141, std=1488.1344
4.  feature\_maps里面的内容多了relu部分，因为是在代码的最后update进去的，所以这部分在feature\_maps的最后：
    
        ......
        vgg_19/conv1/relu1_1 (1, 512, 512, 64)
        vgg_19/conv1/relu1_2 (1, 512, 512, 64)
        vgg_19/conv2/relu2_1 (1, 256, 256, 128)
        vgg_19/conv2/relu2_2 (1, 256, 256, 128)
        vgg_19/conv3/relu3_1 (1, 128, 128, 256)
        vgg_19/conv3/relu3_2 (1, 128, 128, 256)
        vgg_19/conv3/relu3_3 (1, 128, 128, 256)
        vgg_19/conv3/relu3_4 (1, 128, 128, 256)
        ......
5.  其他输出变量已经检查，与第一种方式没有差异，说明魔改的结果是正确的。

## 3. 使用NewCheckpointReader并自定义模型 ##

这种方法分两部分说明。  
第一部分简单说明如何从预训练模型中拿到权重参数；第二部分详细说明如何将预训练权重系数赋值给新定义的模型，并进行测试。

下面是第一部分的代码：

# -*- coding: utf-8 -*-
    from tensorflow.python import pywrap_tensorflow as wrap
    
    
    def main():
        model_path = r'E:\pretrained_model\tf1x\vgg_19.ckpt'
        reader = wrap.NewCheckpointReader(model_path)
    
        variables_shape = reader.get_variable_to_shape_map()
        variables_dtype = reader.get_variable_to_dtype_map()
        for key in variables_shape.keys():
            print(key, variables_shape.get(key), variables_dtype.get(key))
    
        print('\n')
        print(reader.has_tensor("vgg_19/mean_rgb"))
        rgb_mean = reader.get_tensor("vgg_19/mean_rgb")
        print(rgb_mean)
    
    
    if __name__ == '__main__':
        main()

上述代码有几个需说明的地方：

1.  NewCheckpointReader用于加载预训练模型的权重
2.  get\_variable\_to\_shape\_map()和get\_variable\_to\_dtype\_map()可以查看权重参数的shape和dtype
3.  get\_tensor()可以获得权重参数，返回的是`numpy`数组

打印出来的结果如下，其中有两个比较巧妙的参数，`global_step`和`vgg_19/mean_rgb`，mean\_rgb打印出来有具体数值：

global_step [] <dtype: 'int64'>
    vgg_19/conv2/conv2_2/biases [128] <dtype: 'float32'>
    vgg_19/conv2/conv2_2/weights [3, 3, 128, 128] <dtype: 'float32'>
    vgg_19/conv1/conv1_1/biases [64] <dtype: 'float32'>
    vgg_19/conv1/conv1_1/weights [3, 3, 3, 64] <dtype: 'float32'>
    vgg_19/conv1/conv1_2/biases [64] <dtype: 'float32'>
    vgg_19/conv1/conv1_2/weights [3, 3, 64, 64] <dtype: 'float32'>
    ......
    vgg_19/mean_rgb [3] <dtype: 'float32'>
    ......
    vgg_19/fc8/weights [1, 1, 4096, 1000] <dtype: 'float32'>
    
    
    [123.68 116.78 103.94]

下面是第二部分的代码：

# -*- coding: utf-8 -*-
    import os
    import cv2
    import tensorflow as tf
    import numpy as np
    from tensorflow.python import pywrap_tensorflow as wrap
    
    os.environ['CUDA_VISIBLE_DEVICES'] = "/gpu:0"
    slim = tf.contrib.slim
    
    
    def vgg19(inputs,
              scope_name='vgg_19'):
        with tf.variable_scope(scope_name):
    
            net_config = [
                [64, 2],
                [128, 2],
                [256, 4],
                [512, 4],
                [512, 4],
            ]  # [filters, blocks]
    
            feature_maps = { }
            x = inputs
            for i, config in enumerate(net_config):
                filters = config[0]
                for j in range(config[1]):
                    conv_name = 'conv%d_%d' % (i + 1, j + 1)
                    relu_name = 'relu%d_%d' % (i + 1, j + 1)
    
                    x = tf.layers.conv2d(x, filters, [3, 3],
                                         padding='same',
                                         name=conv_name)
                    feat_map_name = x.op.name.replace('/BiasAdd', '')
                    feature_maps[feat_map_name] = x
    
                    x = tf.nn.relu(x, name=relu_name)
                    feature_maps[x.op.name] = x
    
                x = tf.layers.max_pooling2d(x, (2, 2), (2, 2),
                                            name='pool%d' % (i + 1))
                feat_map_name = x.op.name.replace('/MaxPool', '')
                feature_maps[feat_map_name] = x
    
            return x, feature_maps
    
    
    def visualize_feature_map(feature_map,
                              col_nums=None,
                              gap_value=0.5,
                              gap_width=10,
                              gap_height=10):
        # 见本文档前面
        pass
    
    
    def _get_pretrained_tensor_name(name):
        block_num = int(name.split('/')[1][4:].split('_')[0])
        name = name.replace('vgg_19', 'vgg_19/conv%d' % block_num)
        name = name.replace('kernel', 'weights').replace('bias', 'biases')
        return name
    
    
    def main():
        image_file = r'E:\images\lena512color.tiff'
        model_path = r'E:\pretrained_model\tf1x\vgg_19.ckpt'
        inputs_ = tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
        outputs, feature_map_dict = vgg19(inputs_)
        trainable_vars = tf.trainable_variables()
    
        # use NewCheckpointReader to get weights
        reader = wrap.NewCheckpointReader(model_path)
    
        sess = tf.Session()
        sess.run(tf.global_variables_initializer())
    
        # print trainable variables before assignment
        for var in trainable_vars:
            print(var)
            print(sess.run(var)[:, :, 0, 0])
            break
    
        # trainable variables assignment
        print('\n')
        for i, var in enumerate(trainable_vars):
            name = _get_pretrained_tensor_name(var.op.name)
            sess.run(var.assign(reader.get_tensor(name)))
    
        # print trainable variables after assignment
        for var in trainable_vars:
            print(var)
            name = _get_pretrained_tensor_name(var.op.name)
            print(sess.run(var)[:, :, 0, 0])
            print('pretrained weight:')
            print(reader.get_tensor(name)[:, :, 0, 0])
            break
    
        # test case
        inputs = cv2.imread(image_file)
        inputs = np.expand_dims(inputs, axis=0)
        out, feature_maps = sess.run([outputs, feature_map_dict],
                                     feed_dict={ 
                                         inputs_: inputs,
                                     })
    
        # print shape of feature maps
        print('\n')
        for key in feature_maps.keys():
            print(key, feature_maps.get(key).shape)
    
        feature_map = feature_maps.get('vgg_19/relu3_4')
        feature_map = np.squeeze(feature_map)
        image = visualize_feature_map(feature_map)
        image = np.clip(image * 255, 0, 255).astype(np.uint8)
        cv2.imwrite('lena_feature_map_vgg_conv3_4--2.png', image)
    
        # print statistics for feature map
        print('\n')
        for i in range(5):
            mean_val = np.mean(feature_map[:, :, i])
            std = np.std(feature_map[:, :, i])
            min_val = np.min(feature_map[:, :, i])
            max_val = np.max(feature_map[:, :, i])
            print(i + 1, " min=%.4f, max=%.4f, mean=%.4f, std=%.4f" % (
                min_val, max_val, mean_val, std))
    
    
    if __name__ == '__main__':
        main()

比如针对vgg loss这个需求，通常我们不需要最后的全连接层，所以本着节省算力及显存的目的，上述新定义的模型将全连接部分给去掉了，并且feature map / variables 的名字也重新进行了定义，这种情况下就无法使用restore功能加载预训练参数，只能使用赋值的方式。

上述代码整体分为两个部分，一是权重参数赋值，二是跟之前一样的测试用例。

下面对赋值的流程进行说明：

1.  使用placeholder创建计算图并获取`trainable_vars`
2.  使用`NewCheckpointReader`加载预训练模型的权重参数
3.  创建Session并初始化全局变量
4.  使用`var.assign()`方法给权重参数赋值

上述代码打印出来的结果如下：

<tf.Variable 'vgg_19/conv1_1/kernel:0' shape=(3, 3, 3, 64) dtype=float32_ref>
    [[ 0.04975817 -0.0374901  -0.04425776]
     [ 0.03555809  0.08642714  0.05649987]
     [-0.07783681 -0.03184588 -0.07609541]]
    （sess.run(tf.global_variables_initializer())之后打印了kernel的一部分，为随机初始化的结果）
    
    <tf.Variable 'vgg_19/conv1_1/kernel:0' shape=(3, 3, 3, 64) dtype=float32_ref>
    [[ 0.39416704  0.37740308 -0.04594866]
     [ 0.2671299   0.09986369 -0.34100872]
     [-0.07573577 -0.2803425  -0.41602272]]
    pretrained weight:
    [[ 0.39416704  0.37740308 -0.04594866]
     [ 0.2671299   0.09986369 -0.34100872]
     [-0.07573577 -0.2803425  -0.41602272]]
     （权重参数赋值之后，有一次打印了kernel的一部分，同时也打印了预训练模型中对应的部分，可以看到kernel被成功赋值）
    
    
    vgg_19/conv1_1 (1, 512, 512, 64)
    vgg_19/relu1_1 (1, 512, 512, 64)
    vgg_19/conv1_2 (1, 512, 512, 64)
    vgg_19/relu1_2 (1, 512, 512, 64)
    vgg_19/pool1 (1, 256, 256, 64)
    ......
    vgg_19/conv5_4 (1, 32, 32, 512)
    vgg_19/relu5_4 (1, 32, 32, 512)
    vgg_19/pool5 (1, 16, 16, 512)
    （检查featuremap的名字和shape）
    
    
    1  min=0.0000, max=9201.5811, mean=386.3745, std=737.3252
    2  min=0.0000, max=7389.5913, mean=1412.0540, std=616.6437
    3  min=0.0000, max=3323.7239, mean=400.2662, std=522.4063
    4  min=0.0000, max=4319.3765, mean=369.9904, std=644.4222
    5  min=0.0000, max=8997.2305, mean=905.1512, std=1288.8953
    （打印 relu3_4，并与之前的两种方法对比数值，结果是一样的，说明整体流程没什么问题）

最后来看一下代码中保存的那张feature map图：  
![在这里插入图片描述][watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA5ouc6Ziz_size_20_color_FFFFFF_t_70_g_se_x_16_pic_center]

[https_github.com_tensorflow_models_tree_r1.13.0_research_slim]: https://github.com/tensorflow/models/tree/r1.13.0/research/slim
[http_download.tensorflow.org_models_vgg_19_2016_08_28.tar.gz]: http://download.tensorflow.org/models/vgg_19_2016_08_28.tar.gz
[watermark_type_ZHJvaWRzYW5zZmFsbGJhY2s_shadow_50_text_Q1NETiBA5ouc6Ziz_size_20_color_FFFFFF_t_70_g_se_x_16_pic_center]: /images/20220829/fa9a9368d3354f9ba8e9e5a0c4fc3c06.png