python – Tensorflow中的多维聚集

2023年2月28日 228次阅读

这个问题的一般解决方案正在
this github问题上进行,但我想知道是否有使用tf.gather(或其他东西)来使用多索引实现数组索引的变通方法.我想出的一个解决方案是将多idx中的每个索引与张量形状的累积乘积相乘,产生适合索引展平张量的索引：

import tensorflow as tf
import numpy as np

def __cumprod(l):
    # Get the length and make a copy
    ll = len(l)
    l = [v for v in l]

    # Reverse cumulative product
    for i in range(ll-1):
        l[ll-i-2] *= l[ll-i-1]

    return l

def ravel_multi_index(tensor, multi_idx):
    """
    Returns a tensor suitable for use as the index
    on a gather operation on argument tensor.
    """

    if not isinstance(tensor, (tf.Variable, tf.Tensor)):
        raise TypeError('tensor should be a tf.Variable')

    if not isinstance(multi_idx, list):
        multi_idx = [multi_idx]

    # Shape of the tensor in ints
    shape = [i.value for i in tensor.get_shape()]

    if len(shape) != len(multi_idx):
        raise ValueError("Tensor rank is different "
                        "from the multi_idx length.")

    # Work out the shape of each tensor in the multi_idx
    idx_shape = [tuple(j.value for j in i.get_shape()) for i in multi_idx]
    # Ensure that each multi_idx tensor is length 1
    assert all(len(i) == 1 for i in idx_shape)

    # Create a list of reshaped indices. New shape will be
    # [1, 1, dim[0], 1] for the 3rd index in multi_idx
    # for example.
    reshaped_idx = [tf.reshape(idx, [1 if i !=j else dim[0]
                    for j in range(len(shape))])
                for i, (idx, dim)
                in enumerate(zip(multi_idx, idx_shape))]

    # Figure out the base indices for each dimension
    base = __cumprod(shape)

    # Now multiply base indices by each reshaped index
    # to produce the flat index
    return (sum(b*s for b, s in zip(base[1:], reshaped_idx[:-1]))
        + reshaped_idx[-1])

# Shape and slice starts and sizes
shape = (Z, Y, X) = 4, 5, 6
Z0, Y0, X0 = 1, 1, 1
ZS, YS, XS = 3, 3, 4

# Numpy matrix and index
M = np.random.random(size=shape)
idx = [
    np.arange(Z0, Z0+ZS).reshape(ZS,1,1),
    np.arange(Y0, Y0+YS).reshape(1,YS,1),
    np.arange(X0, X0+XS).reshape(1,1,XS),
]

# Tensorflow matrix and indices
TM = tf.Variable(M)
TF_flat_idx = ravel_multi_index(TM, [
    tf.range(Z0, Z0+ZS),
    tf.range(Y0, Y0+YS),
    tf.range(X0, X0+XS)])
TF_data = tf.gather(tf.reshape(TM,[-1]), TF_flat_idx)

with tf.Session() as S:
    S.run(tf.initialize_all_variables())

    # Obtain data via flat indexing
    data = S.run(TF_data)

    # Check that it agrees with data obtained
    # by numpy smart indexing
    assert np.all(data == M[idx])

然而,由于this(当前)限制限制广播到等级3的张量,这仅适用于等级3的张量.

目前我只能想到进行链式收集,转置,收集,转置,收集,但这不太可能有效.例如

shape = (8, 9, 10)
A = tf.random_normal(shape)
data = tf.gather(tf.transpose(tf.gather(A, [1, 3]), [1,0,2]), ...)

有任何想法吗？

最佳答案 听起来你想要
gather_nd.

点赞