beginh
beginh

Reputation: 1153

How to apply mask to a tensor and keep its original shape

I have two tensors: one containing data and the other mask of boolean values. I would like to set all values in data tensor to zero, if boolean values are False, while keeping the original shape of data tensor. So far I can achieve it only while mask is a numpy array.

Since https://www.tensorflow.org/api_docs/python/tf/boolean_mask influences shape of the tensor, I cannot use it.

How to do that?

import numpy as np
import tensorflow as tf
tf.enable_eager_execution()

# create dummy data
data_np = np.ones((4,2,3))
mask_np = np.array([[True, True],[False, True],[True, True],[False, False]])

# prepare tensors
data = tf.convert_to_tensor(data_np)
mask = tf.convert_to_tensor(mask_np)

# how to perform the same while avoiding numpy?
mask = np.expand_dims(mask, -1)
data *= mask

Upvotes: 4

Views: 9963

Answers (2)

Vikram Baliga
Vikram Baliga

Reputation: 504

From the tf.boolean_mask documentation, you can find this:

See also: tf.ragged.boolean_mask, which can be applied to both dense and ragged tensors, and can be used if you need to preserve the masked dimensions of tensor (rather than flattening them, as tf.boolean_mask does).

Use it to retain shapes as indicated in the official documentation. Do not use any other/extra methods. Find the usage details here : Ragged Boolean Mask

Upvotes: 0

Vlad
Vlad

Reputation: 8585

Use tf.cast() and tf.expand_dims():

import tensorflow as tf
import numpy as np

mask_np = np.array([[True, True],[False, True],[True, True],[False, False]])
data_np = np.ones((4,2,3))

mask = tf.convert_to_tensor(mask_np, dtype=tf.bool)
mask = tf.expand_dims(tf.cast(mask, dtype=tf.float32), axis=len(mask.shape))
data = tf.convert_to_tensor(data_np, dtype=tf.float32)

result = mask * data

print(result.numpy())
# [[[1. 1. 1.]
#   [1. 1. 1.]]
# 
#  [[0. 0. 0.]
#   [1. 1. 1.]]
# 
#  [[1. 1. 1.]
#   [1. 1. 1.]]
# 
#  [[0. 0. 0.]
#   [0. 0. 0.]]]

Upvotes: 4

Related Questions