how to apply a mask to a tensor and maintain its original shape… here is a solution to the problem.
how to apply a mask to a tensor and maintain its original shape
I have two tensors: one containing the data and the other containing the bool value mask. If the bool value is False, I want to set all values in the data tensor to zero while keeping the original shape of the data tensor.
So far, I’ve only been able to implement mask when it’s a numpy array.
自 https://www.tensorflow.org/api_docs/python/tf/boolean_mask Affects the shape of the tensor, I can’t use it.
How?
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
Solution
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.]]]