|
Декоратор tf.function для ускорения выполнения функцийНа предыдущем занятии мы с вами построили двухслойную полносвязную НС для классификации изображений цифр. И написали свою процедуру обучения этой сети средствами Tensorflow. Но один важный нюанс остался в стороне – это перенос вычислений на уровень графов для ускорения расчетов. В Tensorflow любой набор операций можно перевести в графовое представление и выполнять более эффективно. Первый вопрос, который здесь может возникнуть, что такое графы в Tensorflow? В двух словах – это представление процесса вычислений в виде последовательности распределенных операций tf.Operation с тензорами tf.Tensor:
Подробнее о них можно почитать на странице официальной документации: https://www.tensorflow.org/guide/intro_to_graphs Не следует путать эту структуру с вычислительным графом в автоматическом дифференцировании. Это разные вещи. Здесь подразумевается, что любой фрагмент программы, написанный с использованием Tensorflow, можно представить в виде такого графа операций. Например, в программе обучения НС есть ресурсоемкий фрагмент расчета производных и изменения обучаемых параметров: with tf.GradientTape() as tape: f_loss = cross_entropy(y_batch, model(x_batch)) grads = tape.gradient(f_loss, model.trainable_variables) opt.apply_gradients(zip(grads, model.trainable_variables)) Этот фрагмент целиком написан с использованием Tensorflow, поэтому его можно выполнять не с помощью интерпретатора языка Python, а через граф операций. И это будет значительно быстрее. Чтобы перенести выполнение этих строчек на уровень графа Tensorflow, достаточно вынести их в функцию, у которой прописать декоратор @tf.function: @tf.function def train_batch(x_batch, y_batch): with tf.GradientTape() as tape: f_loss = cross_entropy(y_batch, model(x_batch)) grads = tape.gradient(f_loss, model.trainable_variables) opt.apply_gradients(zip(grads, model.trainable_variables)) return f_loss А, затем, в цикле обучения вызвать эту функцию: for n in range(EPOCHS): loss = 0 for x_batch, y_batch in train_dataset: loss += train_batch(x_batch, y_batch) print(loss.numpy()) После запуска программы увидим заметное ускорение процесса обучения. Ускорение вычислений в таких графах происходит за счет:
Кроме того, граф Tensorflow можно использовать в других средах без интерпретатора Python, например, в мобильных устройствах, используя язык Java. Введение в tf.functionДавайте теперь детальнее изучим порядок использования и возможности функции tf.function(). Как мы уже с вами отмечали, он автоматически преобразовывает программу на Tensorflow в граф. И это может быть любая программа, не обязательно связанная с обучением НС и вычислением градиентов. Например, можно определить некую функцию с матричными вычислениями: def function_tf(x, y): s = tf.zeros_like(x, dtype=tf.float32) s = s + tf.matmul(x, y) for n in range(10): s = s + tf.matmul(s, y) * x return s А, затем вызвать стандартным способом и через граф. Чтобы оценивать скорость их работы, я записал следующий декоратор: def test_function(fn): def wrapper(*args, **kwargs): start = time.time() fn(*args, **kwargs) dt = time.time() - start print(f"Время обработки: {dt} сек") return wrapper Сформировал две матрицы 1000х1000 элементов: SIZE = 1000 x = tf.ones((SIZE, SIZE), dtype=tf.float32) y = tf.ones_like(x, dtype=tf.float32) Определил представление функции через граф с помощью функции tf.function(): function_tf_graph = tf.function(function_tf) И вызвал их, используя свой декоратор для оценки скорости их работы: test_function(function_tf)(x, y) test_function(function_tf_graph)(x, y) В итоге, получился следующий результат: Время обработки:
0.12494015693664551 сек
Этот пример показывает, что использование графа не всегда приводит к выигрышу в скорости вычислений. Дело в том, что при первом вызове function_tf_graph() Tensorflow строит граф операций и только потом приступает к вычислениям. Суммарное время на построение графа и выполнение оказывается больше, чем простой вызов исходной функции. Но, что если мы предполагаем вызов функции function_tf() в цикле, то есть, много раз. Для тестирования такого поведения добавим вдекораторtest_function() цикл: def test_function(fn): def wrapper(*args, **kwargs): start = time.time() for n in range(10): fn(*args, **kwargs) dt = time.time() - start print(f"Время обработки: {dt} сек") return wrapper Теперь, при оценке скорости работы функций, получаем результат: Время обработки:
1.0309865474700928 сек
Как видим, скорость выполнения в графе выше, чем напрямую интерпретатором языка Python. Мало того, если функция function_tf(), к которой применяется декоратор tf.function, вызывает какие-либо другие функции, то к ним также автоматически применяется этот декоратор. Например: def function_for(s, x, y): for n in range(10): s = s + tf.matmul(s, y) * x return s def function_tf(x, y): s = tf.zeros_like(x, dtype=tf.float32) s = s + tf.matmul(x, y) return function_for(s, x, y) В итоге, весь этот фрагмент программы будет преобразован в граф, а затем, выполнен. И из этих примеров вы можете заметить, что в граф успешно преобразовываются не только тензоры и функции пакета Tensorflow, но и конструкции языка Python, такие как циклы, условные операторы, генераторы последовательностей (range) и т. п. В целом, Tensorflow корректно позволяет транслировать Python-программу в графовое представление. Однако оставит в программе только то, что влияет на конечный результат вычислений. Например, функция print() будет здесь отброшена после первого вызова функции: def function_tf(x, y): print("вызов функции print") s = tf.zeros_like(x, dtype=tf.float32) s = s + tf.matmul(x, y) return function_for(s, x, y) Смотрите, при прямом вызове функции function_tf() мы видим десять сообщений, а при вызове через function_tf_graph() – только одно. Это, как раз, связано с тем, что при первом вызове function_tf_graph() Tensorflow выполняет «трассировку» кода программы в графовое представление и выполняет ее полностью. А при последующих вызовах работает только граф, в котором остались операции, нужные исключительно для получения конечных результатов вычислений. И функция print() здесь явно лишняя. Более подробную информацию об ограничениях транслирования Python-программ в графы, можно почитать в официальной документации: https://www.tensorflow.org/guide/function#limitations Общие рекомендации здесь следующие:
Надеюсь, теперь вы имеете представление о функции tf.function и эквивалентном декораторе, а также понимаете когда и как его применять для ускорения работы программы. Видео по теме |
