tensorflow output_data[output_id == i] = input_data[input_id == i]代码分析

以你之姓@ 2022-02-01 12:43 203阅读 0赞

import tensorflow as tf
    import os
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    
    with tf.Graph().as_default(), tf.Session() as sess:
        input_data = tf.constant([0.1, 0.2, 0.3])
        input_ids = tf.constant([3, 1, 6])
        output_data = tf.constant([0., 0., 0., 0., 0.])
        output_ids = tf.constant([6, 3, 1, 3, 0])
        # From TF v1.13
        #s = tf.argsort(input_ids)
        # Before TF v1.13
        s = tf.contrib.framework.argsort(input_ids)
        input_ids_s = tf.gather(input_ids, s)
        n = tf.size(input_ids)
        output_idx_s = tf.minimum(tf.searchsorted(input_ids_s, output_ids), n - 1)
        searchsorted = tf.searchsorted(input_ids_s, output_ids), n - 1
        gather1 = tf.gather(input_ids_s, output_idx_s)
        equal = tf.equal(output_ids, tf.gather(input_ids_s, output_idx_s))
        gather2 = tf.gather(s, output_idx_s)
        gather3 = tf.gather(input_data, tf.gather(s, output_idx_s))
        output_data = tf.where(tf.equal(output_ids, tf.gather(input_ids_s, output_idx_s)),
                               tf.gather(input_data, tf.gather(s, output_idx_s)),
                               output_data)
        print("output_data",sess.run(output_data))
        print("s", sess.run(s))
        print("input_ids_s", sess.run(input_ids_s))
        print("n", sess.run(n))
        print("searchsorted", sess.run(searchsorted))
        print("input_idx_s",sess.run(output_idx_s))
        print("gather1",sess.run(gather1))
        print("equal", sess.run(equal))
        print("gather2",sess.run(gather2))
        print("gather3",sess.run(gather3))
        print("out_idx_s", sess.run(output_idx_s))

这个代码的功能就是很简单的numpy代码：output\_data\[output\_id == i\] = input\_data\[input\_id == i\],来自于：

[https://stackoverflow.com/questions/56073800/output-dataoutput-ids-i-input-datainput-ids-i-in-tensorflow][https_stackoverflow.com_questions_56073800_output-dataoutput-ids-i-input-datainput-ids-i-in-tensorflow]

只用高端api难以发论文的，因为你用的都是已存在的成熟的高阶代码，哪怕是排列组合也早都有人用过了。

所以tf.core还是要学。你见过用keras写新的随机梯度下降的吗？有的tensorflow都不好使，人家都用c++写。

如果是工业界，那就用越成熟的api越好，毕竟越稳定，越快出成果，越出钱钱啊。

输出：前面是输入，为了看的更明白。

input_data = tf.constant([0.1, 0.2, 0.3])
        input_ids = tf.constant([3, 1, 6])
        output_data = tf.constant([0., 0., 0., 0., 0.])
        output_ids = tf.constant([6, 3, 1, 3, 0])
    
    
    output_data [0.3 0.1 0.2 0.1 0. ]
    s [1 0 2]
    input_ids_s [1 3 6]
    n 3
    searchsorted (array([2, 1, 0, 1, 0], dtype=int32), 2)
    input_idx_s [2 1 0 1 0]
    gather1 [6 3 1 3 1]
    equal [ True  True  True  True False]
    gather2 [2 0 1 0 1]
    gather3 [0.3 0.1 0.2 0.1 0.2]
    out_idx_s [2 1 0 1 0]

大概的思路是这样的，首先找到output\_ind在输入中的索引位置，然后用这个索引gather在input\_data上取数。

如果input\_data是事先根据input\_ids排好序的，那就简单了。只需要两句代码。

还有一个问题要注意，两个ind不一定一模一样，所以还要用equal判断下是否存在于output\_ind

s是input\_id的排序索引

input\_ids\_s 是input\_id 的排序

n是input\_ids的尺寸

tf.searchsort就是在一个已经排序好的tensor中搜索另一个tensor的位置信息。所以这里返回的是tensor的位置。

tf.minimum是为了

[https_stackoverflow.com_questions_56073800_output-dataoutput-ids-i-input-datainput-ids-i-in-tensorflow]: https://stackoverflow.com/questions/56073800/output-dataoutput-ids-i-input-datainput-ids-i-in-tensorflow