tf.boolean_mask 使用方法

import tensorflow as tf
import numpy as np

a = np.concatenate([np.random.rand(2,5,4),np.random.randint(0,2,(2,5,1))],axis=-1)
mask = a[...,-1]>0



sess = tf.Session()
print(a)
print(mask)
print('result:\n')
print(sess.run(tf.boolean_mask(a, mask))) #[1 3]

[[[0.83543017 0.24411466 0.60749395 0.16342278 1.        ]

  [0.69746243 0.61816951 0.70101376 0.25052667 0.        ]

  [0.93887323 0.50630576 0.22931419 0.88845506 1.        ]

  [0.07509089 0.71023028 0.78258055 0.96150005 1.        ]

  [0.61848273 0.95956269 0.77063229 0.36396516 0.        ]]

 [[0.27858061 0.60960442 0.35739185 0.41887888 1.        ]

  [0.44177715 0.38762254 0.50739743 0.82358793 0.        ]

  [0.57369092 0.98604878 0.12832417 0.2824307  0.        ]

  [0.4353096  0.27706318 0.1393903  0.37290264 0.        ]

  [0.32199883 0.07746477 0.13626271 0.64940313 0.        ]]]

[[ True False  True  True False]

 [ True False False False False]]

result:

[[0.83543017 0.24411466 0.60749395 0.16342278 1.        ]

 [0.93887323 0.50630576 0.22931419 0.88845506 1.        ]

 [0.07509089 0.71023028 0.78258055 0.96150005 1.        ]

 [0.27858061 0.60960442 0.35739185 0.41887888 1.        ]]