如何将自定义层放置在内置的预训练模型中？

Question

我们正在尝试在预训练的 imagenet 模型中添加自定义层。对于顺序或非顺序模型，我们可以轻松做到这一点。但这里有一些要求。

首先，我们不想透露整个 imagenet 模型并处理所需的内部层。假设对于DenseNet，我们需要以下层，并进一步获取它们的输出形状以连接一些自定义层。

vision_model = tf.keras.applications.DenseNet121(
    input_shape=(224,224,3),
    include_top = False, 
     weights='imagenet')

for i, layer in enumerate(vision_model.layers):
    if layer.name in ['conv3_block12_concat', 'conv4_block24_concat']:
        print(i,'\t',layer.trainable,'\t  :',layer.name)

        if layer.name == 'conv3_block12_concat':
            print(layer.get_output_shape_at(0)[1:])  # (28, 28, 512)

        if layer.name == 'conv4_block24_concat':
            print(layer.get_output_shape_at(0)[1:])  # (14, 14, 1024)

整个需求可以如下展示

绿色指示灯基本上是密集网的过渡层。

在上图中，密集网络模型具有（假设）5 blocks，其中，我们要选择block 3 和block 4，并添加一些自定义层，然后将它们合并以引导最终输出。

此外，DenseNet 的块（块 1 到 5）应尽可能公开其预训练的 imagenet 权重。我们喜欢在需要时控制冻结和解冻预训练层。

我们如何使用tf.keras 有效地实现？或者，如果您认为有更好的方法来做同样的事情，请提出建议。

---

比方说，自定义块是这样的

class MLPBlock(tf.keras.layers.Layer):
    def __init__(self, kernel_num=32, kernel_size=(3,3), strides=(1,1), padding='same'):
        super(ConvModule, self).__init__()
        # conv layer
        self.conv = tf.keras.layers.Conv2D(kernel_num, 
                        kernel_size=kernel_size, 
                        strides=strides, padding=padding)
        # batch norm layer
        self.bn   = tf.keras.layers.BatchNormalization()

    def call(self, input_tensor, training=False):
        x = self.conv(input_tensor)
        x = self.bn(x, training=training)
        return tf.nn.relu(x)

动机

我正在尝试实施this 文书工作，他们做了这样的事情。最初，这篇论文是免费的，但现在不是了。但下面是他们方法的主要框图。

Answer 1

我无法访问该论文，所以我只是构建了一个示例，就像您绘制的那样：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, models

class ConvBlock(layers.Layer):
    def __init__(self, kernel_num=32, kernel_size=(3,3), strides=(1,1), padding='same'):
        super(ConvBlock, self).__init__()
        # conv layer
        self.conv = layers.Conv2D(kernel_num, 
                        kernel_size=kernel_size, 
                        strides=strides, padding=padding)
        # batch norm layer
        self.bn = layers.BatchNormalization()

    def call(self, input_tensor, training=False):
        x = self.conv(input_tensor)
        x = self.bn(x, training=training)
        return tf.nn.relu(x)

vision_model = keras.applications.DenseNet121(
    input_shape=(224,224,3),
    include_top = False,
    weights='imagenet')

# Control freeze and unfreeze over blocks
def set_freeze(block, unfreeze):
    for layer in block:
        layer.trainable = unfreeze

block_1 = vision_model.layers[:7]
block_2 = vision_model.layers[7:53]
block_3 = vision_model.layers[53:141]
block_4 = vision_model.layers[141:313]
block_5 = vision_model.layers[313:]

set_freeze(block_1, unfreeze=False)
set_freeze(block_2, unfreeze=False)

for i, layer in enumerate(vision_model.layers):
    print(i,'\t',layer.trainable,'\t  :',layer.name)

layer_names = ['conv3_block12_concat', 'conv4_block24_concat', 'conv5_block16_concat']
vision_model_outputs = [vision_model.get_layer(name).output for name in layer_names]

custom_0 = ConvBlock()(vision_model_outputs[0])
custom_1 = ConvBlock()(layers.UpSampling2D(2)(vision_model_outputs[1]))
cat_layer = layers.concatenate([custom_0, custom_1])

last_conv_num = 2
custom_2 = layers.UpSampling2D(4)(vision_model_outputs[2])
outputs = layers.concatenate([ConvBlock()(cat_layer) for i in range(last_conv_num)] + [custom_2])
model = models.Model(vision_model.input, outputs)

keras.utils.plot_model(model, "./Model_structure.png", show_shapes=True)

运行代码，你会看到block1和block2被冻结了，

因为完整模型的情节很长，所以我只发布了几个sn-p：