CODEWITHSUNDEEP
pythontensorflow
Gaddy
Gaddy

Reputation: 463

Nested tf.map_fn with slow performance

I run the code below in order to get a padding matrix from a given matrix of indices (words_chars_ids shape is (6,200,20)). The result has the shape (6,200,20,emb_size) where for each entry in the output it holds a tensor of ones or zeroes (with size emb_size).

I have two questions:

  1. Is there a more elegant way to implement this method (without the nested map_fn)

  2. The performance seems to be quite slow - is there a more efficient way to achieve the result ?

def get_padding_mask(words_chars_ids, emb_size):

    padding_mask = tf.map_fn(
        lambda x: tf.map_fn(
            lambda y: tf.map_fn(
                lambda z: tf.cond(tf.less(z, 1),
                                  lambda: tf.zeros([emb_size, ], dtype=tf.int32),
                                  lambda: tf.ones([emb_size, ], dtype=tf.int32)
                                  ),
                y),
            x),
        words_chars_ids)
    return padding_mask

Upvotes: 1

Views: 148

Answers (1)

javidcf
javidcf

Reputation: 59681

You can do the same thing simply as:

def get_padding_mask(words_chars_ids, emb_size):
    mask = tf.dtypes.cast(words_chars_ids >= 1, tf.int32)
    return tf.tile(tf.expand_dims(mask, -1), [1, 1, 1, emb_size])

Upvotes: 1

Related Questions