Эффективный метод для нахождения индекса первого появления числа в пакетных данных, отличных от цикла for

Question

Я делаю задачу, в которой у меня есть данные в виде кадров, хранящихся в пакетах.Размер пакета подобен (batch_size, 400), я хочу найти индекс первого появления числа 1 в каждом кадре длиной 400.

В настоящее время я использую для цикла цикл по размеру, но поскольку данные очень большие, это занимает очень много времени

Любой другой эффективный метод, использующий некоторую матричную операцию в tenorflow или numpy,

Answer 1

В TensorFlow:

import tensorflow as tf

def index_of_first_tf(batch, value):
    eq = tf.equal(batch, value)
    has_value = tf.reduce_any(eq, axis=-1)
    _, idx = tf.math.top_k(tf.cast(eq, tf.int8))
    idx = tf.squeeze(idx, -1)
    return tf.where(has_value, idx, -tf.ones_like(idx))

В NumPy:

import numpy as np

def index_of_first_np(batch, value):
    eq = np.equal(batch, value)
    has_value = np.any(eq, axis=-1)
    idx = np.argmax(eq, axis=-1)
    idx[~has_value] = -1
    return idx

Тесты:

import tensorflow as tf

batch = [[0, 1, 2, 3],
         [1, 2, 1, 0],
         [0, 2, 3, 4]]
value = 1

print(index_of_first_np(batch, value))
# [ 1  0 -1]

with tf.Graph().as_default(), tf.Session() as sess:
    print(sess.run(index_of_first_tf(batch, value)))
    # [ 1  0 -1]