python - Theano/numpy 高级索引

Question

我有一个 4d theano 张量(例如形状为 (1, 700, 16, 95000))和一个形状为 (1, 700, 16, 1024) 的 4d 'mask' 张量，使得mask 是我需要原始张量的索引。如何使用掩码来索引张量？像sample[mask]或sample[:,:,:,mask]这样的东西似乎并没有真正起作用。

我还尝试使用二进制掩码，但由于张量相当大，我收到“设备内存不足”异常。

关于如何从张量获取索引的其他想法也将非常感激。

谢谢

Answer 1

因此，在缺乏答案的情况下，我决定使用计算量更大的解决方案，该解决方案将我的数据展开为索引张量，向索引添加偏移量以将它们带到全局位置，对数据进行索引并将其 reshape 回原来的样子。

我在此处添加我的测试代码，包括矩阵的(注释掉的)解决方案。

def theano_convertion(els, inds, offsets):
  els = T.flatten(els)
  inds = T.flatten(inds) + offsets
  return T.reshape(els[inds], (2, 3, 16, 5))


if __name__ == '__main__':
  # command: np.transpose(t[range(2), indices])
  # t = np.random.randint(0, 10, (2, 20))
  # indices = np.random.randint(0, 10, (5, 2))

  t = np.random.randint(0, 10, (2, 3, 16, 20)).astype('int32')
  indices = np.random.randint(0, 10, (2, 3, 16, 5)).astype('int32')
  offsets = np.asarray(range(1, 2 * 3 * 16 + 1), dtype='int32')
  offsets = (offsets * 20) - 20
  offsets = np.repeat(offsets, 5)

  offsets_tens = T.ivector('offsets')
  inds_tens = T.itensor4('inds')
  t_tens = T.itensor4('t')

  func = theano.function(
    [t_tens, inds_tens, offsets_tens],
    [theano_convertion(t_tens, inds_tens, offsets_tens)]
  )

  shaped_elements = []
  flattened_elements = []

  [tmp] = func(t, indices, offsets)

  for i in range(2):
    for j in range(3):
      for k in range(16):
          shaped_elements.append(t[i, j, k, indices[i, j, k, :]])
          flattened_elements.append(tmp[i, j, k, :])

          print shaped_elements[-1] == flattened_elements[-1]