首页 > 解决方案 > 形状为 b*n*3 的张量流张量 T1。形状为 b*n 的 T2 -> 一个布尔张量,指示要在 T1 中接收哪些行

问题描述

我正在尝试在n点 3D 坐标和b批次中提取特定行。基本上我的张量T1是 shape b*n*3。我有另一个布尔T2形状的张量,指示需要采取b * n哪些行。n基本上我的输出应该是b*?*3因为T2每批中可以有不同数量的 1。

我已经使用布尔掩码实现了以下内容,但输出与预期不同,输出形状为(?,)但不是(b*?*3)

# expand T2 to (b,n,3). i.e. 0 replicates to (0,0,0) and so is 1

mask = tf.tile(tf.expand_dims(T2,2), [1,1,3])

# query using boolean mask where there are 1s

valid_KPs = tf.boolean_mask(T1, tf.cast(mask, tf.int32))

标签: pythontensorflow

解决方案


由于每个示例的选定元素的数量可能不同,因此不能表示为适当的张量。一种选择是使用参差不齐的张量。他们不能做普通张量所能做的一切,但可以实现你想要的,例如:

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    # Input data
    t1 = tf.constant([
        [
            [ 1,  2,  3],
            [ 4,  5,  6],
            [ 7,  8,  9],
        ],
        [
            [10, 11, 12],
            [13, 14, 15],
            [16, 17, 18],
        ],
    ])
    t2 = tf.constant([
        [1, 0, 1],
        [0, 1, 0],
    ])
    # Count the number of ones for each row in T2
    c = tf.reduce_sum(t2, axis=1)
    # Ragged ranges for each row
    r = tf.ragged.range(c)
    # Sorting indices so indices with a one are first
    s = tf.argsort(t2, axis=1, direction='DESCENDING', stable=True)
    # First axis dimension index
    idx0 = tf.expand_dims(tf.range(tf.shape(t1)[0]), 1) * tf.ones_like(r)
    # 2D index for getting indices of ones on each row
    idx_s = tf.stack([idx0, r], axis=-1)
    # Get indices of ones
    idx1 = tf.gather_nd(s, idx_s)
    # 2D index to get indices of selected vectors in T1
    idx = tf.stack([idx0, idx1], axis=-1)
    # Get selected vectors
    result = tf.gather_nd(t1, idx)
    # Print result
    print(sess.run(result))
    # <tf.RaggedTensorValue [[[1, 2, 3], [7, 8, 9]], [[13, 14, 15]]]>

推荐阅读