Transforming keras model output during training and use multiple losses

Question

       if self.use_embedding:
            input_x = Input(shape=(self.input_length,), name='embedding_input')
            x = Embedding(self.input_dim, self.embedding_dim)(input_x)
        else:
            input_x = Input(shape=(self.input_length,self.input_dim), name='notes_input')
            x = input_x

        encoder_input_list = [input_x]
        encoded = self._build_encoder(x)
        self.encoder = Model(inputs=encoder_input_list, outputs=encoded)

        
        encoded_input = Input(shape=(self.latent_rep_size,), name='encoded_input')
        
        if self.use_embedding:
            input_decoder_x = Input(shape=(self.output_dim,), name='embedding_input_decoder_start')
            #decoder_x = Embedding(self.output_dim, self.output_dim, input_length=1)(input_decoder_x)
            decoder_x = input_decoder_x
        else:
            input_decoder_x = Input(shape=(self.output_dim,), name='input_decoder_start')
            decoder_x = input_decoder_x

        autoencoder_decoder_input_list = [input_decoder_x, encoded]
        decoder_input_list = [input_decoder_x, encoded_input]
        autoencoder_input_list = [input_x, input_decoder_x]
        autoencoder_output_list = []

        
        if self.teacher_force:
            ground_truth_input = Input(shape=(self.output_length, self.output_dim), name='ground_truth_input')
            decoder_input_list.append(ground_truth_input)
            autoencoder_decoder_input_list.append(ground_truth_input)
            autoencoder_input_list.append(ground_truth_input)
        else:
            ground_truth_input = None

        if self.history:
            history_input = Input(shape=(self.latent_rep_size,), name='history_input')
            decoder_input_list.append(history_input)
            autoencoder_decoder_input_list.append(history_input)
            autoencoder_input_list.append(history_input)
        else:
            history_input = None

        decoded= self._build_decoder(decoder_x, encoded_input, ground_truth_input, history_input)
        
        loss_list = []
        loss_weights_list = []
        sample_weight_modes = []

        loss_weights_list.append(1.0)
        sample_weight_modes.append('temporal')
        loss_list.append(self.vae_loss)
        metrics_list = ['accuracy']
        
        decoder_output = decoded
        
        self.decoder = Model(inputs=decoder_input_list, outputs=decoder_output, name='decoder')

        decoder_final_output = self.decoder(autoencoder_decoder_input_list)

        if isinstance(decoder_final_output, list):
            autoencoder_output_list.extend(decoder_final_output)
        else:
            autoencoder_output_list.append(decoder_final_output)
            
        self.autoencoder = Model(inputs=autoencoder_input_list, outputs=autoencoder_output_list, name='autoencoder')
        self.autoencoder.compile(optimizer=self.optimizer,
                                 loss=loss_list,
                                 loss_weights=loss_weights_list,
                                 sample_weight_mode=sample_weight_modes,
                                 metrics=metrics_list)

Here is my code for a variational sequence to sequence model. I want to transform the decoder output during the training so that It gives me two other matrices and use two different loss functions for the new matrices. So the total loss should be like L_total = L_mat1 + L_mat2 + L_mat3 - KL loss. Is there a way to do it in Keras? I have the ground truth values for the transformed matrices. Any good soul would help me with this?

Answer 1

I've added an example on how to implement a matrix split operation to this colab notebook: https://colab.research.google.com/drive/127DY9OUWasFQzM9G2AH4RQO8ryhSTJny?usp=sharing

The short answer is that you can use a lambda to split the output of any layer.

In my simple example:

  inputs = layers.Input(shape=(2,))
  d = layers.Dense(10)(inputs)
  out1, out2 = layers.Lambda(lambda x: (x[:, :5], x[:, 5:]))(d)
  model = keras.Model(inputs, [out1, out2])  
  model.compile(loss=['mse', 'mae'])
  return model

The lambda layer splits the first 5 columns of the matrix into out1 and the last 5 cols into out2.