Issues with masking in custom model

TensorFlow General Discussion
1

I am trying to generate a model that takes as an input a tuple consisting of a rank-1 tensor of id’s and a rank-2 tensor with feature vector’s corresponding to those ID’s. With this input, I would like to use the ID’s to run the corresponding feature vectors through a ID-defined feed-forward neural net. Each feed-forward neural net gives a single output, which are then collated into a rank-1 tensor of the same shape as ID’s that is further processed. As a side-note, each instance of an ID-defined feed-forward neural net should carry the same weights.

I have tried the following, but run into some sort of issue with masking and the if statement that I am having trouble diagnosing.

class SAPAYOA(tf.Model):
    def __init__(self, species):
        super().__init__()
        #These are all feed forward neural networks.
        self.ANN = [NN_H,NN_C,NN_N,NN_N,NN_S]
        self.species = species
    def call(self,inputs):      
        species, aev = inputs
        assert species.shape = aev.shape[:-1]
        species_ = tf.reshape(species, [tf.size(species)])
        aev = tf.reshape(aev, [tf.size(species),48])
        
        #This doesn't work, but batch size dependent response is also a bit underwhelming.
        output = tf.zeros(species.get_shape().as_list(), dtype = aev.dtype)
         
        for i,m in enumerate(self.ANN):         
            filtered_species = tf.keras.layers.Masking(mask_value=i)(species)            
            X_indices = tf.cast(tf.where(~filtered_species._keras_mask),tf.int32)
            #This if statement is problematic
            if X_indices.shape[0] > 0: 
                X_vectors = tf.gather_nd(indices=X_indices, params=aev)
                output += tf.scatter_nd(X_indices,m(X_vectors),output.shape)
        return output

Any help would be appreciated!

3

Hi @Andres_Cordoba,

Sorry for the delay in response.
I’d recommend to create a mask directly with tf.equal(species, i) to select the relevant indices instead of using tf.keras.layers.Masking layer and use tf.where(mask) to gather the corresponding feature vectors with tf.gather_nd. In addition, instead of scatter_nd, use tf.tensor_scatter_nd_add to add the processed outputs back into the output tensor at the correct indices and ensure shapes of species and aev are aligned properly. Here’s updated code snippet for your reference.

class SAPAYOA(tf.Model):
    def __init__(self, species):
        super().__init__()
        # These are all feed-forward neural networks.
        self.ANN = [NN_H, NN_C, NN_N, NN_N, NN_S]  # Assuming NN_H, NN_C, etc. are defined elsewhere
        self.species = species

    def call(self, inputs):
        species, aev = inputs
        assert species.shape == aev.shape[:-1]
        
        batch_size = tf.shape(species)[0]
        aev = tf.reshape(aev, [batch_size, -1])  # Assuming 48 features per ID
        
        # Initialize output tensor with zeros (same shape as species)
        output = tf.zeros_like(species, dtype=aev.dtype)
        
        for i, m in enumerate(self.ANN):         
            mask = tf.equal(species, i)  
            X_indices = tf.where(mask) 
            
            if tf.size(X_indices) > 0:
                X_vectors = tf.gather_nd(aev, X_indices)
                m_output = m(X_vectors)
                output = tf.tensor_scatter_nd_add(output, X_indices, m_output)
        
        return output

Hope this helps.Thank You.