Nhập trục theo chỉ mục trong Tensorflow

Question

Tôi có hàng chục đầu vào 3D [batch, n_classes - k, 5]và hàng chục chỉ số 2D [batch, n_classes - k]. kcó thể ở trong [0, n_classes)Như một ví dụ, hãy lấy n_classes=3, k=1:

X = tf.constant([
    [[0.36636186, 0.45606998, 0.785176  , 0.19967379, 0. ],
     [0.2799339 , 0.9548653 , 0.7378969 , 0.5543541 , 1. ]],

    [[0.07455064, 0.9868869 , 0.77224475, 0.19871569, 0. ],
     [0.19579114, 0.0693613 , 0.100778  , 0.01822183, 1. ]],

    [[0.684233  , 0.4401525 , 0.12203824, 0.4951769 , 0. ],
     [0.47417384, 0.09783416, 0.49161586, 0.47347176, 0. ]]
])

idcs = tf.constant([
    [0, 2],
    [0, 1],
    [1, 2]
])

Trong đó các phần tử trong idcslà các giá trị lớp (chỉ số). Tôi đang cố gắng áp dụng Xdọc theo trục 1 cho các chỉ số bị thiếu trong phạm vi (0, n_classes)bằng cách đặt chúng thành vectơ số không, tức là,

tf.constant([
    [[0.36636186, 0.45606998, 0.785176  , 0.19967379, 0. ],
     [0.        , 0.        , 0.        , 0.        , 0. ],  # missing 1 in `idcs`
     [0.2799339 , 0.9548653 , 0.7378969 , 0.5543541 , 1. ]],

    [[0.07455064, 0.9868869 , 0.77224475, 0.19871569, 0. ],
     [0.19579114, 0.0693613 , 0.100778  , 0.01822183, 1. ],
     [0.        , 0.        , 0.        , 0.        , 0. ]], # missing 2 in `idcs`

    [[0.        , 0.        , 0.        , 0.        , 0. ],  # missing 0 in `idcs`
     [0.684233  , 0.4401525 , 0.12203824, 0.4951769 , 0. ],
     [0.47417384, 0.09783416, 0.49161586, 0.47347176, 0. ]]
])

Tôi không chắc làm thế nào để diễn đạt điều này trong tensorflow. Tôi đã cân nhắc việc tạo một tensor số không [batch, n_classes, 5]và gán các chỉ số hiện tại Xdọc theo trục 1, nhưng không được phép gán trong tensor. Có cách nào dễ dàng để thực hiện điều này trong tensorflow không?

Ví dụ, nếu tôi muốn diễn đạt điều này bằng Numpy, tôi có thể thử một cái gì đó như:

X = np.array([
    [[0.36636186, 0.45606998, 0.785176  , 0.19967379, 0. ],
     [0.2799339 , 0.9548653 , 0.7378969 , 0.5543541 , 1. ]],

    [[0.07455064, 0.9868869 , 0.77224475, 0.19871569, 0. ],
     [0.19579114, 0.0693613 , 0.100778  , 0.01822183, 1. ]],

    [[0.684233  , 0.4401525 , 0.12203824, 0.4951769 , 0. ],
     [0.47417384, 0.09783416, 0.49161586, 0.47347176, 0. ]]
])

idcs = np.array([
    [0, 2],
    [0, 1],
    [1, 2]
])

n_classes = 3
batch_size = 3

# selectors
x = np.repeat(np.arange(idcs.shape[0]), 2)  # [0, 0, 1, 1, 2, 2]
y = idcs.ravel()  # [0, 2, 0, 1, 1, 2]

z = np.zeros((batch_size, n_classes, 5))
z[x, y] = np.reshape(X, [x.shape[0], 5])
z

# array([[[0.36636186, 0.45606998, 0.785176  , 0.19967379, 0.        ],
#         [0.        , 0.        , 0.        , 0.        , 0.        ],
#         [0.2799339 , 0.9548653 , 0.7378969 , 0.5543541 , 1.        ]],
# 
#        [[0.07455064, 0.9868869 , 0.77224475, 0.19871569, 0.        ],
#         [0.19579114, 0.0693613 , 0.100778  , 0.01822183, 1.        ],
#         [0.        , 0.        , 0.        , 0.        , 0.        ]],
# 
#        [[0.        , 0.        , 0.        , 0.        , 0.        ],
#         [0.684233  , 0.4401525 , 0.12203824, 0.4951769 , 0.        ],
#         [0.47417384, 0.09783416, 0.49161586, 0.47347176, 0.        ]]])

Answer 1

Tôi đã giải quyết vấn đề này với một chút hack bằng cách chuyển đổi tensor thành thưa thớt, và sau đó ngay lập tức trở lại dày đặc:

batch_size, n_inputs, _ = X.shape.as_list()
n_classes = 3

sparse_indices = tf.concat([
        tf.reshape(tf.repeat(tf.range(batch_size, dtype=tf.int64), n_inputs * 5), [-1, 1]),
        tf.reshape(tf.repeat(idcs, 5), [-1, 1]),
        tf.reshape(tf.tile(tf.range(5, dtype=tf.int64), [n_inputs * batch_size]), [-1, 1]),
    ],
    axis=1
)

# ravel X to 1d, create a sparse tensor for non-zero indices and then
# expand back to dense as a hack for filling in the zeros
X_ravel = tf.reshape(X, shape=[-1])
tf.sparse.to_dense(
    tf.sparse.SparseTensor(
        sparse_indices,
        X_ravel,
        dense_shape=[batch_size, n_classes, 5],
    ),
)

Như mong đợi, mang lại:

<tf.Tensor: shape=(3, 3, 5), dtype=float32, numpy=
array([[[0.36636186, 0.45606998, 0.785176  , 0.19967379, 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.2799339 , 0.9548653 , 0.7378969 , 0.5543541 , 1.        ]],

       [[0.07455064, 0.9868869 , 0.77224475, 0.19871569, 0.        ],
        [0.19579114, 0.0693613 , 0.100778  , 0.01822183, 1.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ]],

       [[0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.47417384, 0.09783416, 0.49161586, 0.47347176, 0.        ],
        [0.684233  , 0.4401525 , 0.12203824, 0.4951769 , 0.        ]]],
      dtype=float32)>