合并两个训练有素的网络以按顺序进行推理

Question

我正在尝试合并两个经过训练的神经网络。 我有两个经过训练的 Keras 模型文件 A 和 B。

模型 A 用于图像超分辨率，模型 B 用于图像着色。

我正在尝试合并两个经过训练的网络，以便更快地推断 SR+着色。 （我不愿意使用一个网络来完成 SR 和着色任务。我需要使用两个不同的网络来完成 SR 和着色任务。）

关于如何合并两个 Keras 神经网络的任何提示？

Answer 1

只要网络 A 的输出形状与模型 B 的输入形状兼容，就有可能。

由于 tf.keras.models.Model 继承自 tf.keras.layers.Layer，因此您可以使用 Model，就像在创建 keras 模型时使用 Layer 一样。

---

一个简单的例子：

让我们首先创建 2 个简单的网络 A 和 B，其约束条件是 B 的输入与 A 的输出具有相同的形状。

import tensorflow as tf 
A = tf.keras.models.Sequential(
    [
        tf.keras.Input((10,)),
        tf.keras.layers.Dense(5, activation="tanh")
    ],
    name="A"
)
B = tf.keras.models.Sequential(
    [
        tf.keras.Input((5,)),
        tf.keras.layers.Dense(10, activation="tanh")
    ],
    name="B"
)

然后我们可以将这两个模型合并为一个，在这种情况下使用功能 API（这也完全可以使用 Sequential API）：

merged_input = tf.keras.Input((10,))
x = A(merged_input)
merged_output = B(x)
merged_model = tf.keras.Model(inputs=merged_input, outputs=merged_output, name="merged_AB")

导致以下网络：

>>> merged_model.summary()
Model: "merged_AB"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_3 (InputLayer)         [(None, 10)]              0         
_________________________________________________________________
A (Sequential)               (None, 5)                 55        
_________________________________________________________________
B (Sequential)               (None, 10)                60        
=================================================================
Total params: 115
Trainable params: 115
Non-trainable params: 0
_________________________________________________________________