有条件地在TF模型中拆分执行

Question

我正在按两个部分/块制作张量流模型。输入+主体只是几个卷积而已。在我当前的模型中，卷积块进入一个完全连接的块，并做出决定。

我想尝试做的是创建多个完全连接的块，并使用一个处理卷积块输出的层。最大输出索引将对应于完全连接的块的索引。卷积块会将其输出提供给选定的决策者。

这个想法有点来自多任务深度神经网络（https://arxiv.org/abs/1901.11504），但是在选择要使用的头部时采用了更多的手动方法。

我的问题是这样：

我需要有条件地更改模型的执行，但我不确定如何做到这一点。我现在正在使用Keras，并试图使其与Model和Sequential一起使用，这让我很困惑。我对API的主要经验是使用model.train（...），看来我可能需要采取更手动的方法？

提前谢谢！我是神经网络的新手，所以很抱歉这是一个愚蠢的问题。

编辑：模型决定由左边的方框决定。在此示例中，我是说这是一个完全连接的层，然后进行softmax激活。

• 输入：卷积的展平输出
• Output：完全连接的层的最终输出的argmax。 argmax将成为决策者之一的指数

实际上，我正计划使用其他方法，但为简单起见，这就是我刚才所描述的。

Answer 1

要做您想做的事情有点复杂。但是仍然可能。

from tensorflow.keras import layers, models

height = 64
width = 64
n_channels = 3
n_layers = 4

# We have the simple convolutional output here
inp = layers.Input(shape=(height, width, n_channels))
conv_out = layers.Conv2D(32, (3,3), padding='same')(inp)
conv_out = layers.Flatten()(conv_out)

# (Model decision) Your model producing the correct ID of the fully connected layer (batch size, 4)
out = layers.Dense(n_layers, activation='softmax')(conv_out)

# This is the tricky bit. We are creating indices to gather the parameters from the corresponding layer
# Both best_out, range_out are (batch_size, 1) concatenated to (batch_size, 2)
# They would be like [(0, 2), (1, 3), (2, 0), (3, 0), .... (batch_size-1, 2)] type indices
best_out = layers.Lambda(lambda x: tf.expand_dims(tf.argmax(x, axis=-1), axis=-1))(out)
range_out = layers.Lambda(lambda x: tf.math.cumsum(tf.ones_like(x), exclusive=True))(best_out)
gather_ids = layers.Concatenate(axis=-1)([range_out, best_out])

# Compute the output for all dense layers
final_outputs = [layers.Dense(10)(conv_out) for _ in range(n_layers)]
# Stack the outputs (batch size, n_layers, 10)
final_output = layers.Lambda(lambda x: tf.stack(x, axis=1))(final_outputs)
# Gather the outputs corresponding to correct indices (batch_size, 10)
final_out = layers.Lambda(lambda x: tf.gather_nd(x[0], x[1]))([final_output, gather_ids])

model = models.Model(inputs=inp, outputs=[best_out, final_out])

model.summary()

简单测试

让我们看看如果给定索引传入一个数组，是否可以为给定索引获得相同的向量。

x0 = np.ones(shape=(3, height, width, n_channels))
x0[2,:,:,:] = np.random.normal(size=(height, width, n_channels))
y0 = model.predict(x0)

for ri in range(y0[0].size):
  print(y0[0][ri])
  print(y0[1][ri])
  print('\n')

输出（3个元素的批处理）

Max index => [1]
[-0.3227892   1.1366667  -0.43643862  1.0002458  -1.2778764  -1.1994138
  0.18452626  0.89942855  0.33907077 -0.22196433]


Max index => [1]
[-0.3227892   1.1366667  -0.43643862  1.0002458  -1.2778764  -1.1994138
  0.18452626  0.89942855  0.33907077 -0.22196433]


Max index => [2]
[ 0.18891329  0.18289518  0.43710774 -0.2164327  -0.03203449  0.16306376
 -0.0089941   0.3098799  -1.0531532  -0.34144163]