Question

如果我想获取每3个元素的总和，该怎么办？

test_arr = [1,2,3,4,5,6,7,8]

听起来像是地图功能

map_fn(arr, parallel_iterations = True, lambda a,b,c : a+b+c)

并且map_fn(test_arr)的结果应为

[6,9,12,15,18,21]

等于

[(1+2+3),(2+3+4),(3+4+5),(4+5+6),(5+6+7),(6+7+8)]

在查看了官方文档https://www.tensorflow.org/api_docs/python/tf/map_fn

之后，我已经找到了解决方案

import tensorflow as tf

def tf_map_elements_every(n, tf_op, input, dtype):
    if n >= input.shape[0]:
        return tf_op(input)
    else:
        return tf.map_fn(
            lambda params: tf_op(params),
            [input[i:i-n+1] if i !=n-1 else input[i:] for i in range(n)], 
            dtype=dtype
        )

测试

t = tf.constant([1, 2, 3, 4, 5, 6, 7, 8])
op = tf_map_elements_every(3, tf.reduce_sum, t, tf.int32)

sess = tf.Session()
sess.run(op)

[Out]: array([ 6, 9, 12, 15, 18, 21])

Answer 1

更容易：使用列表理解。将列表切成3个元素的片段，然后将每个片段相加。将它们包装在列表中。

[sum(test_arr[i-2:i+1]) 
    for i in range(2, len(test_arr))]

Answer 2

简单地循环遍历数组，直到最后三岁为止。

# Takes a collection as argument
def map_function(array):
    # Initialise results and i
    results = []
    int i = 0

    # While i is less than 3 positions away from the end of the array
    while(i <= (len(array) - 3)):
        # Add the sum of the next 3 elements in results
        results.append(array[i] + array[i + 1] + array[i + 2]

        # Increment i
        i += 1

    # Return the array
    return results

Tensorflow：如何在每n个元素上并行执行一个操作？

2 个答案: