INNER CODE UNIT · Python
weighted_expert_output
drawbridge/keras-mmoe · mmoe.py:191
weighted_expert_output = expert_outputs * K.repeat_elements(expanded_gate_output, self.units, axis=1)
final_outputs.append(K.sum(weighted_expert_output, axis=2))
return final_outputs
def compute_output_shape(self, input_shape):
"""
Method for computing the output shape of the MMoE layer.
:param input_shape: Shape tuple (tuple of integers)
:return: List of input shape tuple where the size of the list is equal to the number of tasks
"""
assert input_shape is not None and len(input_shape) >= 2
output_shape = list(input_shape)
output_shape[-1] = self.units
output_shape = tuple(output_shape)