tf.gather_nd

参考

tf.gather_nd - 云+社区 - 腾讯云


 

将params中的切片收集到一个由指标指定形状的张量中。

tf.gather_nd(
    params,
    indices,
    batch_dims=0,
    name=None
)

指标是一个k维整数张量，最好考虑为(K-1)张量的指标到帕拉姆，其中每个元素定义了帕拉姆的一个切片:

output[\\(i_0, ..., i_{K-2}\\)] = params[indices[\\(i_0, ..., i_{K-2}\\)]]

而在tf.gather索引中，将切片定义为params的第一个维度，而在tf.gather_nd中，索引将切片定义为params的第一个N个维度，其中N = indices.shape[-1]。指标的最后一个维度最多可以是参数的秩:

indices.shape[-1] <= params.rank

指标的最后一个维度对应于元素(if指标)。shape[-1] == parameters .rank或slice (if索引)。形状[-1]< params.rank)沿维度指标。形状参数的[1]。输出张量是有形状的：

indices.shape[:-1] + params.shape[indices.shape[-1]:]

此外，“params”和“indices”都可以有M个完全匹配的领先批处理维度。在本例中，'batch_dims'必须是M。注意，在CPU上，如果发现一个out of bound索引，将返回一个错误。在GPU上，如果发现一个out of bound索引，则在相应的输出值中存储一个0。下面的一些例子。简单的索引成一个矩阵:

    indices = [[0, 0], [1, 1]]
    params = [['a', 'b'], ['c', 'd']]
    output = ['a', 'd']

切片索引成一个矩阵:

    indices = [[1], [0]]
    params = [['a', 'b'], ['c', 'd']]
    output = [['c', 'd'], ['a', 'b']]

索引成一个3张量:

    indices = [[1]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[['a1', 'b1'], ['c1', 'd1']]]
 
 
    indices = [[0, 1], [1, 0]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['c0', 'd0'], ['a1', 'b1']]
 
 
    indices = [[0, 0, 1], [1, 0, 1]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = ['b0', 'b1']

下面的例子适用于只有索引具有领先额外维度的情况。如果“params”和“indexes”都具有领先的批处理维度，则使用“batch_dims”参数以批处理模式运行gather_nd。批量索引成一个矩阵:

    indices = [[[0, 0]], [[0, 1]]]
    params = [['a', 'b'], ['c', 'd']]
    output = [['a'], ['b']]

成批切片索引成矩阵:

    indices = [[[1]], [[0]]]
    params = [['a', 'b'], ['c', 'd']]
    output = [[['c', 'd']], [['a', 'b']]]

批量索引成一个3张量:

    indices = [[[1]], [[0]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[[['a1', 'b1'], ['c1', 'd1']]],
              [[['a0', 'b0'], ['c0', 'd0']]]]
 
    indices = [[[0, 1], [1, 0]], [[0, 0], [1, 1]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[['c0', 'd0'], ['a1', 'b1']],
              [['a0', 'b0'], ['c1', 'd1']]]
 
 
    indices = [[[0, 0, 1], [1, 0, 1]], [[0, 1, 1], [1, 1, 0]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['b0', 'b1'], ['d0', 'c1']]

批量使用“参数”和“索引”的例子:

    batch_dims = 1
    indices = [[1], [0]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['c0', 'd0'], ['a1', 'b1']]
 
    batch_dims = 1
    indices = [[[1]], [[0]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [[['c0', 'd0']], [['a1', 'b1']]]
 
    batch_dims = 1
    indices = [[[1, 0]], [[0, 1]]]
    params = [[['a0', 'b0'], ['c0', 'd0']],
              [['a1', 'b1'], ['c1', 'd1']]]
    output = [['c0'], ['b1']]

参数：

• params:一个张量。用来收集值的张量。
• indices:一个张量。必须是下列类型之一:int32、int64。指数张量。
• name:操作的名称(可选)。
• batch_dims:整数或标量“张量”。批量尺寸的数量。

返回值：

• 一个张量。具有与params相同的类型。