Keras functional API Layer name not captured with TimeDistributed wrapper

Question

class Net2:
    @staticmethod
    def build_cat_branch(inputs,category_size):
        x = TimeDistributed(Dense(category_size))(inputs)
        x = Activation('softmax', name="cat_output")(x)
        return x

    @staticmethod
    def build_t_branch(inputs):
        x = TimeDistributed(Dense(1, activation='relu', name="t_output"))(inputs)
        return x

    @staticmethod
    def build_full_model(timestep_len,hidden_size,category_size,num_features,dropout,rec_drop):
        inputs = Input(shape=(timestep_len,num_features),name="Input")
        bn = BatchNormalization()(inputs)
        lstm = LSTM(hidden_size, return_sequences=True, dropout=dropout, recurrent_dropout=rec_drop,name="LSTM")(bn)
        bn2 = BatchNormalization()(lstm)
        cat_branch = Net2.build_cat_branch(bn2,category_size)
        t_branch = Net2.build_t_branch(bn2)
        model = Model(inputs=inputs,outputs=[cat_branch,t_branch],name="Net2")
        return model

When I try to compile this model I get:

ValueError: Unknown entry in loss dictionary: "t_output". Only expected the following keys: ['cat_output', 'time_distributed_2']

(in my model summary, my layer that I name "t_output" has the name "time_distributed_2" instead)

So, basically, my question is what's going on with the TD wrapper that causes the name attribute to not be part of the object returned by the build_t_branch function? Clearly, the "cat_output" name is stored correctly, as the loss dictionary recognizes it, but the output layer I have inside a TimeDistributed wrapper is not saving the user defined layer-name. I know I can get around this by just having all layers defined in a single function without the "branch" functions, but that is besides the point here. Is this a bug in Keras? Any way to get around this without the above mentioned fix?

Answer 1

The reason is you should name the TD instead of Dense. Be cautious with the bracket. Change it to:

x = TimeDistributed(Dense(1, activation='relu'), name="t_output")(inputs)