如何优雅地在numpy中删除不必要的元素？

Question

I have an ndarray of shape [batch_size, seq_len, num_features] . However, some of elements in the end of the sequential dimension is not necessary, and therefore I want to drop them and merge the sequential dimension into the batch dimension. For example, the ndarray a I want to manipulate is

batch_size = 2
seq_len = 3
num_features = 1
a = np.random.randn(batch_size, seq_len, num_features)
mask = np.ones((batch_size, seq_len), dtype=np.bool)
mask[0][1:] = 0
mask[1][2:] = 0
"""
>>> a = [[[-0.3908401 ]
  [ 0.89686512]
  [ 0.07594243]]

 [[-0.12256737]
  [-1.00838131]
  [ 0.56543754]]]
mask=[[ True False False]
 [ True  True False]]
"""

where mask is used to indicate whether the elements in a is useful. I can get what I want using the following code

res = []
for seq, m in zip(a, mask):
    res.append(seq[:sum(m)])
np.concatenate(res, axis=0)
"""
>>>array([[0.08676509],
       [0.47162315],
       [0.98070665]])
"""

I'm wondering if there is a more elegant way to do this in numpy?

Answer 1

不确定这是否是您的要问，但结果看起来还不错

res = a[mask]

Answer 2

Since dimensions related to batch and seq are going to be merged, you could reshape both a and mask to 2D array of shape (batch_size * seq_len, num_features) .

Next, simply filter important samples using boolean index. See the code:

mask2d = mask.reshape(-1) # or mask.ravel()
a2d = a.reshape(-1, num_features)
result = a2d[mask2d]