我正在寻找一种 TensorFlow 方法来实现类似于 Python 的 list.index() 函数。
给定一个矩阵和一个要查找的值，我想知道该值在矩阵的每一行中第一次出现。
例如，
m is a <batch_size, 100> matrix of integers
val = 23
result = [0] * batch_size
for i, row_elems in enumerate(m):
result[i] = row_elems.index(val)
我不能假设“val”在每一行中只出现一次，否则我会使用 tf.argmax(m == val) 来实现它。在我的例子中，重要的是获取 first 出现 'val' 而不是任何的索引。
【问题讨论】：
标签： tensorflow
似乎tf.argmax 的工作方式类似于np.argmax（根据the test），当最大值出现多次时，它将返回第一个索引。
你可以使用tf.argmax(tf.cast(tf.equal(m, val), tf.int32), axis=1) 来得到你想要的。但是，目前tf.argmax 的行为在多次出现最大值的情况下是未定义的。
如果您担心未定义的行为，您可以按照@Igor Tsvetkov 的建议将tf.argmin 应用于tf.where 的返回值。
例如，
# test with tensorflow r1.0
import tensorflow as tf
val = 3
m = tf.placeholder(tf.int32)
m_feed = [[0 , 0, val, 0, val],
[val, 0, val, val, 0],
[0 , val, 0, 0, 0]]
tmp_indices = tf.where(tf.equal(m, val))
result = tf.segment_min(tmp_indices[:, 1], tmp_indices[:, 0])
with tf.Session() as sess:
print(sess.run(result, feed_dict={m: m_feed})) # [2, 0, 1]
请注意，当某些行不包含 val 时，tf.segment_min 将引发 InvalidArgumentError。在您的代码中，row_elems.index(val) 也会在 row_elems 不包含 val 时引发异常。
【讨论】：
argmax 上的 TF 文档明确指出：“请注意，在 tie 的情况下，不保证返回值的身份。”这让我相信你不能像 numpy 那样依赖 argmax 返回第一个值，我怀疑这是因为 GPU 等分布式设备上的非确定性行为。
tf.argmax 的 documentation 确实保证“在身份的情况下返回最小索引。”跨度>
看起来有点难看但有效（假设 m 和 val 都是张量）：
idx = list()
for t in tf.unpack(m, axis=0):
idx.append(tf.reduce_min(tf.where(tf.equal(t, val))))
idx = tf.pack(idx, axis=0)
编辑：
正如Yaroslav Bulatov 提到的，您可以使用tf.map_fn 获得相同的结果：
def index1d(t):
return tf.reduce_min(tf.where(tf.equal(t, val)))
idx = tf.map_fn(index1d, m, dtype=tf.int64)
【讨论】：
map_fn不用拆包就能搞定
这是问题的另一种解决方案，假设每一行都有一个命中。
import tensorflow as tf
val = 3
m = tf.constant([
[0 , 0, val, 0, val],
[val, 0, val, val, 0],
[0 , val, 0, 0, 0]])
# replace all entries in the matrix either with its column index, or out-of-index-number
match_indices = tf.where( # [[5, 5, 2, 5, 4],
tf.equal(val, m), # [0, 5, 2, 3, 5],
x=tf.range(tf.shape(m)[1]) * tf.ones_like(m), # [5, 1, 5, 5, 5]]
y=(tf.shape(m)[1])*tf.ones_like(m))
result = tf.reduce_min(match_indices, axis=1)
with tf.Session() as sess:
print(sess.run(result)) # [2, 0, 1]
【讨论】：
这是一个解决方案，它还考虑了矩阵不包含元素的情况（来自 DeepMind 的 github 存储库的解决方案）
def get_first_occurrence_indices(sequence, eos_idx):
'''
args:
sequence: [batch, length]
eos_idx: scalar
'''
batch_size, maxlen = sequence.get_shape().as_list()
eos_idx = tf.convert_to_tensor(eos_idx)
tensor = tf.concat(
[sequence, tf.tile(eos_idx[None, None], [batch_size, 1])], axis = -1)
index_all_occurrences = tf.where(tf.equal(tensor, eos_idx))
index_all_occurrences = tf.cast(index_all_occurrences, tf.int32)
index_first_occurrences = tf.segment_min(index_all_occurrences[:, 1],
index_all_occurrences[:, 0])
index_first_occurrences.set_shape([batch_size])
index_first_occurrences = tf.minimum(index_first_occurrences + 1, maxlen)
return index_first_occurrences
还有：
import tensorflow as tf
mat = tf.Variable([[1,2,3,4,5], [2,3,4,5,6], [3,4,5,6,7], [0,0,0,0,0]], dtype = tf.int32)
idx = 3
first_occurrences = get_first_occurrence_indices(mat, idx)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
sess.run(first_occurrence) # [3, 2, 1, 5]
【讨论】：