[Python] Error During call() Function of Custom Subclassed Tensorflow MoE Model

I wish to implement a custom Mixture of Experts model subclassing keras.Model. Inside the call() function, model gets weights from another model (gating model) and does arithmetic operation on them. Particularly I am having trouble getting weights.

class MoE(keras.Model):
def __init__(self, gating_model, expert_models):
super(MoE, self).__init__()
self.experts = expert_models
for expert in self.experts:
expert.trainable = False
self.gating = gating_model
def call(self, x):
outputs = tf.concat([tf.expand_dims(expert_model(tf.expand_dims(x, axis=0)), axis=2) for expert_model in expert_models], axis=2)
gating_weights = [layer.get_weights() for layer in gating_model.layers]
gating_weights = tf.reshape(gating_weights, tf.shape(outputs))
return tf.reduce_sum(outputs * gating_weights, dim=2)
def get_config(self):
config = super().get_config()
config.update({"experts": self.experts, "gating": self.gating_model})
return config


Throws:

--> 214 gating_weights = [layer.get_weights() for layer in gating_model.layers]
215 gating_weights = tf.reshape(gating_weights, tf.shape(outputs))

NotImplementedError: Exception encountered when calling MoE.call().

numpy() is only available when eager execution is enabled.

Arguments received by MoE.call():
• x=tf.Tensor


Which is not very self-explanatory as I donot have numpy anywhere. Appreciate if anyone can figure this one out.

Continue reading...