оптимизация производительности графа вывода тензорного потока

Я пытаюсь понять больше о некоторых удивительных результатах, которые я вижу при реализации графика tf. Граф, с которым я работаю, — это просто лес (куча деревьев). Это просто простой график прямого вывода, и ничего связанного с обучением. Я делюсь фрагментами для 2 реализации

фрагмент кода 1:

with tf.name_scope("main"):

    def get_tree_output(offset):
        loop_vars = (offset,)
        leaf_indice = tf.while_loop(cond,
                                    body,
                                    loop_vars,
                                    back_prop=False,
                                    parallel_iterations=1,
                                    name="while_loop")
        tree_score = tf.gather(score_tensor, leaf_indice, name="tree-scores")
        output = tf.add(tree_score, output)

    leaf_indices = tf.map_fn(get_tree_output,
                             tree_offsets_tensor,
                             dtype=INT_TYPE,
                             parallel_iterations=n_trees,
                             back_prop=False,
                             name="tree-scores")

    tree_scores = tf.gather(score_tensor, leaf_indices, name="tree-scores")

    output = tf.reduce_sum(tree_scores, name="sum-output")
    output = tf.sigmoid(output, name="sigmoid-output")

фрагмент кода 2:

with tf.name_scope("main"):
    tree_offsets_tensor = tf.constant(tree_offsets, dtype=INT_TYPE, name="tree_offsets_tensor")
    loop_vars = (tree_offsets_tensor,)
    leaf_indices = tf.while_loop(cond,
                                 body,
                                 loop_vars,
                                 back_prop=False,
                                 parallel_iterations=n_trees,
                                 name="while_loop")

    tree_scores = tf.gather(score_tensor, leaf_indices, name="tree-scores")

    output = tf.reduce_sum(tree_scores, name="sum-output")
    output = tf.sigmoid(output, name="sigmoid-output")

Остальной код точно такой же: постоянные тензоры, переменные, условие и тело цикла while. поток и параллелизм также были одинаковыми в обоих случаях. Фрагмент кода 2: выполнение вывода занимает около 500 микросекунд. Фрагмент кода 1: выполнение вывода занимает около 12 миллисекунд.

Разница в том, что во фрагменте 1 я использую map_fn для работы с tree_offset_tensor, тогда как во фрагменте 2 я избавляюсь от этого map_fn и просто напрямую использую этот тензор, поэтому, как я понимаю, во фрагменте 1 метод get_tree_output вызывается с одним элементом из tree_offset_tensor у нас есть несколько while_loop для каждого отдельного значения смещения, тогда как во фрагменте 2 у нас есть только один while_loop, который просто принимает несколько значений смещения (в основном offset_tensor).

Я также попробовал другой вариант для фрагмента, вместо использования map_fn я пишу написанный от руки цикл for

фрагмент кода 1 (вариант цикла):

output = 0
with tf.name_scope("main"):
    for offset in tree_offsets:
        loop_vars = (offset,)
        leaf_indice = tf.while_loop(cond,
                                    body,
                                    loop_vars,
                                    back_prop=False,
                                    parallel_iterations=1,
                                    name="while_loop")
        tree_score = tf.gather(score_tensor, leaf_indice, name="tree-scores")
        output = tf.add(tree_score, output)

    #leaf_indices = tf.map_fn(get_tree_output,
    #    tree_offsets_tensor, dtype=INT_TYPE,
    #    parallel_iterations=n_trees, back_prop=False,
    #    name="tree-scores")

    #tree_scores = tf.gather(score_tensor, leaf_indices, name="tree-scores")

    #output = tf.reduce_sum(tree_scores, name="sum-output")
    output = tf.sigmoid(output, name="sigmoid-output")

Это дает незначительное улучшение: 9 миллисекунд