Numpy 在批处理数据（而非for循环）中查找数字首次出现索引的有效方法

我正在执行一项任务，在该任务中，我以帧的形式批量存储数据。批次的维度类似于（batch_size，400），我想在每个400长度的帧中找到数字1第一次出现的索引

目前我使用的是批量大小上的for循环，但由于数据非常大，因此非常耗时

在tensorflow或numpy中使用某些矩阵运算的任何其他有效方法在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))

在努比：

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]

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

在努比：

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]