Imputation des axes par index dans Tensorflow

Question

J'ai un tenseur d'entrées 3D [batch, n_classes - k, 5]et un tenseur d'indices 2D [batch, n_classes - k]. kpeut être dans [0, n_classes)À titre d'exemple, prenez 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]
])

Où les éléments dans idcssont des valeurs de classe (indices). J'essaie d'imputer le Xlong de l'axe 1 pour les indices manquants dans la plage (0, n_classes)en les définissant sur des vecteurs de zéros, c'est-à-dire,

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. ]]
])

Je ne sais pas trop comment exprimer cela en tensorflow. J'ai envisagé de créer un tenseur de zéros [batch, n_classes, 5]et d'assigner des indices présents le Xlong de l'axe 1, mais l'affectation n'est pas autorisée dans les tenseurs. Existe-t-il un moyen simple d'accomplir cela dans tensorflow?

Par exemple, si je devais exprimer cela dans Numpy, je pourrais essayer quelque chose comme:

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

J'ai résolu cela avec un peu d'un hack en convertissant le tenseur en clairsemé, puis immédiatement de nouveau en dense:

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],
    ),
)

Ce qui, comme prévu, donne:

<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)>