有时候,我们需要减少预训练模型的输出个数。其实只用一行代码就可以实现。没必要对模型先进行拆分,再进行组装。
比如我们先创建一个二输出的模型
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)