tensorflow2.3+多任务学习MTL保存多个模型方案

esmm模块

以下代码是一个简单的esmm的模型，当我们想保存多个模型的时候，而且线上预测的时候，不想传两个label(label1,label2)进去的时候,处理方式如下

def base_model(inputs,output, variable_scope):
    with tf.compat.v1.variable_scope(variable_scope):
        base_model = tf.keras.Model(
            inputs=inputs,
            outputs=output, name=variable_scope)
    return base_model
# 模型
def gen_model():
    input1 = tf.keras.layers.Input(shape=(2,), dtype=tf.float32, name='ty')
 	
   
    input_layers= Dense(units=16,activation='relu')(input1)
    
    a_preds = tf.keras.layers.Dense(units=1, activation='sigmoid')(input_layers)
    c_preds = tf.keras.layers.Dense(units=1, activation='sigmoid')(input_layers)
    a_model=base_model(inputs=input1 ,output=a_preds , variable_scope='a_model')
     c_model=base_model(inputs=input1 ,output=a_preds , variable_scope='c_model')


    # output layer
    a_preds = tf.keras.layers.Lambda(lambda x: x, name='a')(a_preds)
    b_preds = tf.keras.layers.Multiply(name='b')([a_preds, c_preds])
    b_model=base_model(inputs=input1 ,output=a_preds , variable_scope='b_model')

 	a_label = tf.keras.layers.Input(shape=(1,), dtype=tf.float32, name='a_label')
    b_label = tf.keras.layers.Input(shape=(1,), dtype=tf.float32, name='b_label')
    loss_inputs = [a_label, a_preds, b_label, b_preds]

    def sum_loss(inputs):
        a_true, a_pred, b_true, b_pred = inputs
        a_loss = tf.keras.losses.binary_crossentropy(y_true=a_true, y_pred=a_pred)
        b_loss = tf.keras.losses.binary_crossentropy(y_true=b_true, y_pred=b_pred)
        loss = a_loss + b_loss
        return loss

    # loss layer
    loss_layer = tf.keras.layers.Lambda(lambda x: sum_loss(x), name='loss')(loss_inputs)
    outputs = [a_preds, b_preds, loss_layer]
    loss = [lambda y_true, y_pred: tf.keras.losses.binary_crossentropy(y_true, y_pred),
            lambda y_true, y_pred: tf.keras.losses.binary_crossentropy(y_true, y_pred),
            lambda y_true, y_pred: tf.math.reduce_mean(y_pred)]

    metrics = {'a': tf.keras.metrics.AUC(),
               'b': tf.keras.metrics.AUC()}

    model = tf.keras.Model(inputs=[input1,a_label,b_label],
                           outputs=outputs,
                           name="model")
    initial_learning_rate = 0.15
    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate,
        decay_steps=100000,
        decay_rate=0.8,
        staircase=True)

    model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=lr_schedule),
                  loss=loss,
                  loss_weights=[0, 0, 1], metrics=metrics)

    return  a_model,b_model,c_model,model

save model

save model的时候直接调用上述模块
直接

a_model.save('a') 
b_model.save('b')
c_model.save('c')

注意，这样保存的模型，当a_model预测的时候不用输入(a_label,b_label),b_model、c_model也同理，但是model进行预测的时候，需要灌