I would like to perform the weighted addition of three outputs from different Keras layers such that the weights are trainable. How can I achieve this? I am using tensorflow 2.0 as backend for Keras.


2 Answers 2


I solved the problem using subclassing in keras. The code is shown below:

class Wt_Add(keras.layers.Layer):
def __init__(self, units=1, input_dim=1):
    super(Wt_Add, self).__init__()
    w_init = tf.random_normal_initializer()
    self.w1 = tf.Variable(
        initial_value=w_init(shape=(input_dim, units), dtype="float32"),
    self.w2 = tf.Variable(
        initial_value=w_init(shape=(input_dim, units), dtype="float32"),
    self.w3 = tf.Variable(
        initial_value=w_init(shape=(input_dim, units), dtype="float32"),

def call(self, input1, input2, input3):
    return tf.multiply(input1,self.w1) + tf.multiply(input2, self.w2) + tf.multiply(input3, self.w3)


wt_add = Wt_Add(1,1)
sum_layer = wt_add(input1, input2, input3)

You have the following basic operations on layers:

  • tf.keras.layers.Lambda so you can multiply each of your 3 layers with a simple lambda operation
layer1 =  tf.keras.layers.Lambda(lambda x: x * weight1)(layer1)
layer2 =  tf.keras.layers.Lambda(lambda x: x * weight2)(layer2)
layer3 =  tf.keras.layers.Lambda(lambda x: x * weight3)(layer3)

then there is the tf.keras.layers.Average that allows to average layers:

average_layer = tf.keras.layers.Average()([layer1, layer2, layer3])

It's a bit awkward, I think a weighted average would be the best thing here but it does not seem to be available in Keras yet (as far as I know)

  • $\begingroup$ Thanks. However on saving the model, following error is displayed: "Tried to export a function which references untracked object Tensor("36550:0", shape=(), dtype=resource).TensorFlow objects (e.g. tf.Variable) captured by functions must be tracked by assigning them to an attribute of a tracked object or assigned to an attribute of the main object directly." $\endgroup$ Aug 25, 2020 at 7:19
  • $\begingroup$ In addition, following warning is displayed during training: "The following Variables were used a Lambda layer's call (lambda_10), but are not present in its tracked objects: <tf.Variable 'Variable:0' shape=(1, 1) dtype=float32> It is possible that this is intended behavior, but it is more likely an omission. This is a strong indication that this layer should be formulated as a subclassed Layer rather than a Lambda layer." $\endgroup$ Aug 25, 2020 at 7:30
  • $\begingroup$ looks great - thanks I'll take note of this $\endgroup$ Aug 25, 2020 at 9:57
  • $\begingroup$ Do you need to constrain the three learned weights so they sum to 1? Like using a dirichlet? $\endgroup$
    – bw4sz
    Oct 15, 2020 at 18:30

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.