keras的模型剪枝或减少输出

有时候，我们需要减少预训练模型的输出个数。其实只用一行代码就可以实现。没必要对模型先进行拆分，再进行组装。

比如我们先创建一个二输出的模型

from tensorflow.keras import *
import tensorflow as tf
from tensorflow.keras.layers import *
INPUT=Input((32,32,3))
x=Conv2D(4,2,padding='same')(INPUT)
x=Conv2D(4,2,padding='same')(x)
x1=Conv2D(4,2,padding='same')(x)
x1=Conv2D(4,2,padding='same')(x1)
x2=Conv2D(4,2,padding='same')(x)
x2=Conv2D(4,2,padding='same')(x2)
model=Model(inputs=INPUT,outputs=[x1,x2])

下面来看一下该模型的结构

tf.keras.utils.plot_model(model, show_shapes=True, dpi=50)

现在如果要删除右边的分支，输入下面的代码

model=Model(inputs=model.input,outputs=model.get_layer('conv2d_3').output)

再来看看网络结构

tf.keras.utils.plot_model(model, show_shapes=True, dpi=50)