keras中训练和测试阶段的不同层配置

Question

伙计们，我正在尝试在NN中使用不同的路径进行训练和测试阶段，我用boolean为我的模型创建了一个函数，该函数指示训练或测试阶段，但是当使用keras加载模型时 我在训练配置中没有测试网络。这意味着我得到3个输出而不是1个。

   net = load_model('output/weights.100-0.860.h5',custom_objects={'train':False})

这是受到启发的： https://github.com/aitorzip/Keras-ICNet

    def build_bn(height, width, n_classes, weights_path=None, train=False):
    inp = Input(shape=(height, width, 4))
        .....#some codes


        if train:
        aux_1 = Conv2D(n_classes, 1, activation='softmax', name='sub4_out')(aux_1)
        aux_2 = Conv2D(n_classes, 1, activation='softmax', name='sub24_out')(aux_2)
        model = Model(inputs=inp, outputs=[out, aux_2, aux_1])
        # model.summary()

    else:
        model = Model(inputs=inp, outputs=out)

    if weights_path is not None:
        model.load_weights(weights_path, by_name=True)


    return model