Tensorflow Looping over Sliced ​​Назначение переменной в режиме активного исполнения - PullRequest
Новая
       21

Tensorflow Looping over Sliced ​​Назначение переменной в режиме активного исполнения

1 голос
/ 14 апреля 2020

Для некоторого пользовательского кода мне нужно запустить for-l oop, чтобы динамически создать переменную в Tensorflow 2 (с активным режимом исполнения). (В моем собственном коде для значений, которые я пишу в переменную, потребуются градиенты, поэтому я хочу отслеживать вычисления в for l oop, чтобы я мог получать градиенты из autodiff). Мой код работает, но он невероятно медленный. Фактически, это на несколько порядков медленнее, чем выполнение той же операции в numpy.

Я выделил проблему и предоставляю фрагмент кода игрушки, который освещает проблему. Исправление позволит мне исправить мой пользовательский код. 

import numpy as np
import tensorflow as tf
import timeit

N = int(1e5)
data = np.random.randn(N)
def numpy_func(data):
    new_data = np.zeros_like(data)
    for i in range(len(data)):
        new_data[i] = data[i]
    return new_data

def tf_func(data):
    new_data = tf.Variable(tf.zeros_like(data))
    for i in range(len(data)):
        new_data[i].assign(data[i])
    return new_data    

%timeit numpy_func(data)
%timeit tf_func(data)

Ключевые выводы этого фрагмента кода: for-l oop Мне нужно только обновить фрагмент переменной. Срез, который нужно обновить, отличается на каждой итерации. Данные, используемые для обновления, отличаются на каждой итерации (в моем пользовательском коде это результат нескольких простых вычислений, которые зависят от кусочков переменной, здесь я просто использую фиксированный массив для изоляции проблемы.)

Я использую Tensorflow 2, и в идеале код тензорного потока должен запускаться с активным выполнением, так как части пользовательских операций зависят от исполнения.

Я новичок в Tensorflow и буду очень признателен за помощь в решении этой проблемы.

Большое спасибо, Макс

1 Ответ

0 голосов
/ 14 апреля 2020

TensorFlow никогда не будет очень быстрым при таком использовании. Идеальным решением было бы векторизовать ваши вычисления, чтобы не требовалось явно указывать l oop, но это зависит от того, что именно вы вычисляете (вы можете опубликовать еще один вопрос об этом, если хотите). Тем не менее, вы можете получить немного лучшую производительность, используя tf.function. Я немного изменил вашу функцию, чтобы иметь new_data в качестве выходного параметра, так как tf.function не позволяет вам создавать переменные после первого вызова (но на самом деле, если вы удалите параметр new_data, это будет также работает, так как tf.function найдет переменную в глобальной области видимости). 

import numpy as np
import tensorflow as tf
import timeit

# Input data
N = int(1e3)
data = np.random.randn(N)

# NumPy
def numpy_func(data, new_data):
    new_data[:] = 0
    for i in range(len(data)):
        new_data[i] = data[i]

new_data = np.zeros_like(data)
%timeit numpy_func(data, new_data)
# 143 µs ± 4.41 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# TensorFlow
def tf_func(data, new_data):
    new_data.assign(tf.zeros_like(data))
    for i in range(len(data)):
        new_data[i].assign(data[i])
new_data = tf.Variable(tf.zeros_like(data))
%timeit tf_func(data, new_data)
# 119 ms ± 3.68 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# tf.function
# This is equivalent to using it as a decorator
tf_func2 = tf.function(tf_func)
new_data = tf.Variable(tf.zeros_like(data))
tf_func2(data, new_data)  # First call is slower
%timeit tf_func2(data, new_data)
# 3.55 ms ± 40.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Это было выполнено на CPU, результаты могут значительно отличаться на GPU. В любом случае, как вы можете видеть, с tf.function он все еще более чем в 20 раз медленнее, чем NumPy, но также более чем в 30 раз быстрее, чем функция Python.

...