Напомним, что графы вычислений — представление структуры данных в виде графа. Думаю, как выглядят математические графы все видели. Узлы — операции: функция активации, сложение, вычитание, ReLU. Ребра — потоки данных, они отражают зависимости между операциями. Обычно ребро направлено от узла-источника к узлу-получателю и передаёт результаты вычислений от одного узла к другому. 

В отличие от Pytorch, где структура данных выстраивается налету после начала обучения нейронки – в TensorFlow граф статичен. Да, мы немного теряем пространства для экспериментов. Но зачем они нам, когда мы пишем нейронку под реальный продакшен?

Автоматическая оптимизация для ленивых

Начнем, пожалуй, с самых простых и полуавтоматизированных способов оптимизации графов вычислений и первый из них – XLA. Высокоуровневый компилятор, ускоряет выполнения операций линейной алгебры — компилирует их в машинный удобочитаемый для вашего GPU или TPU код. 

Работает он относительно просто, для начала производится анализ графа на выявление возможностей для оптимизации. Объединение последовательных операций или слияние вычислительных графов. Зачем преобразования операций или вообще их удаление. И только после этих процедур XLA компилирует полученный граф в код, специализируя его под TPU или GPU. 

Чтобы воспользоваться компилятором, достаточно установить нужный флаг tf.function в функции. 

@tf.function(jit_compile=True)
def model_inference(input):
    return model(input)

Или применить флаг непосредственно перед нужными операциями. Если очень хочется применить компилятор глобально — установите переменную окружения:

import os
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'

Этот способ автоматизированный и далеко не всегда гарантирует "качественную" оптимизацию. Поэтому давайте подумаем о том, как получить оптимальные плюсы от оптимизации и даже уменьшения веса модели другими методами — заморозкой графа и чекпоинтом. 

Заморозка графа или как получить всю карту вычислений в одном файле

Мы мумифицируем/замораживаем переменные и их веса в константы — так мы сможем выгрузить граф в отдельный единый файл. Как и с виртуальной средой разработки или попросту проектом со своими зависимостями, граф вычислений можно приложить к репозиторию. Но для нас самое главное — он нужен для оптимизации вычислений.

Принцип заморозки заключается не только в переводе переменных в константы, но и удаление всех нод графа, которые не нужны для инференса модели. И это преимущество метода. Благодаря ему значительно снизить вес файла. 

Но, безусловно, такую модель нельзя просто так дообучить. 

После того как мы обучили и создали саму модель, нам достаточно удалить узлы обучения и естественно конвертировать все переменные в константы.

# Конвертация переменных модели в константы и удаление узлов обучения

frozen_graph = tf.compat.v1.graph_util.convert_variables_to_constants(
    sess=tf.compat.v1.keras.backend.get_session(),  
  # Получение текущей сессии TensorFlow
    input_graph_def=tf.compat.v1.graph_util.remove_training_nodes(  
      # Удаление узлов обучения
        tf.compat.v1.graph_util.convert_variables_to_constants(  
          # Конвертация переменных в константы
            tf.compat.v1.keras.backend.get_session(),  
          # Текущая сессия TensorFlow
            tf.compat.v1.keras.backend.get_session().graph.as_graph_def(),  
          # Граф модели
            [node.op.name for node in model.outputs]  
          # Имена выходных узлов модели
        )
    ),
    output_node_names=[node.op.name for node in model.outputs]  
  # Имена выходных узлов модели
)

Затем мы сохраняем наш готовый граф в файл

with tf.io.gfile.GFile('frozen_model.pb', 'wb') as f:  
# Открытие файла на запись
    f.write(frozen_graph.SerializeToString())  
# Сериализация замороженного графа и запись его в файл

Чекпоинт — наша точка возрождения в Tensorflow 

Создание легкого чекпоинта в TensorFlow сохраняет состояние модели (веса и оптимизатор) в файл для последующего восстановления и продолжения обучения. Это полезно, если обучение модели занимает много времени, и вы хотите регулярно сохранять ее состояние во время обучения, чтобы избежать потери прогресса в случае сбоя или прерывания. Тот самый cntrl + S, только во фреймворке. 

Ну и зачем нам эти "сохранения"? Дело в том, что чекпоинтами можно управлять. 

И соответственно выжимать все необходимые соки в момент сохранения. 

Например, при помощи модуля tf.train.Saver (...) можно задать число необходимых переменных, которые мы хотим сохранить при инициализации. Туда, например, можно просто вписать все веса. 

С другой стороны, чекпоинты предлагают и другой функционал. Благодаря MetaGraph мы можем сразу строить граф вычислений при загрузке модели из чекпоинта. 

Но мы предлагаем другую идею. Просто продублировать код построения графа из обучения и больше не сохранять метаданные. 

ckpt_filepath = saver.save(sess, filepath, write_meta_graph=False)

Тут мы присваиваем "записи" метаданных False и сохраняем только веса и другие состояние графа при построении чекпоинта. Так можно снизить вес модели чуть ли не в пять раз — пользуемся! 

Но есть еще парочку интересных способов прокачать вашу нейронку и поработать с графами — заняться прунингом, чем-то явно непристойным или квантизацией. 

Прунинг или удаление лишних нейронов в мозгу ИИ

Прунинг или как значительно уменьшить размер модели засчет незначительного понижеия точности. Точнее удаления лишних параметров: весов или нейронов. Например, при помощи оценки градиентов по отношению к параметрам, встроенный Grappler во фреймворк или PruningAPI от google-research убирает самые бесполезные нейроны. 

Реализация от гугла (Pruning) работает по принципу добавления переменных mask, threshold, в которых накапливается информация о весах, дающих верные предсказания. После ненужные веса попросту зануляются и на выходе мы получаем тот самый оптимизированный замороженный граф. 

Например,

pruning_params = {
    'pruning_schedule': sparsity.ConstantSparsity(0.5, begin_step=0, frequency=100) 
  # Определение параметров прунинга
}

pruned_model = sparsity.prune_low_magnitude(model, **pruning_params) 
# Применение прунинга к слоям модели

Сначала создается модель Sequential с несколькими полносвязными слоями. Затем определяются гиперпараметр "разреженности" в виде "расписания"  и частоты.  Эти параметры передаются в функцию prune_low_magnitude, которая применяет прунинг к указанным слоям модели. После этого модель компилируется и обучается так же, как и обычная модель.

Конечно, теоретически вес модели должен уменьшиться, а скорость вырасти.  

Вот только работает ли этот метод для всех типов нейронок? — нет, в конечном счете даже после удаления всех ненужных элементов в модели, у вас остаются эмбендинги, которые прунить мы не советуем. 

Graph Transform Tool и метод квантизации

Подробно о том, как пользоваться Graph Transform tool лучше почитать на Github. 

Дополнительно нужно импортировать утилиту.

import tensorflow as tf
from tensorflow.tools.graph_transforms import TransformGraph

Создайте граф вычислений в TensorFlow и определите список преобразований, которые вы хотите применить к графу. Преобразования могут быть определены в виде списка словарей, каждый из которых содержит имя преобразования и его параметры. 

Например, определим граф для нейронки:

imported_model = tf.saved_model.load(export_dir)
graph_def = imported_model.signatures["serving_default"].graph.as_graph_def()

И прямиком зададим преобразования: 

transforms = [
    {"name": "FoldConstants", "params": {}}
]

optimized_graph_def = TransformGraph(graph_def, input_nodes, output_nodes, transforms)

Применить граф можно так:

optimized_graph_def = TransformGraph(graph_def, input_nodes, output_nodes, transforms)

Здесь мы задаем сам граф, входые и "выходные" ноды, трансформации. 

Да и в целом GTT — мощный инструмент, но не так прост. И далеко не всегда он нужен в разработке, все зависит от ситуации. Назовем его профессиональным режимом оптимизации. 

При помощи инструмента можно проводить квантинизацию — процесс уменьшения точности чисел (обычно из 32-битных float до 8-битных integer). Все просто. У тула есть свои трансформации для этого способа оптимизации. 

transforms = [   
              {'name': 'quantize_weights', 'params': {'dtype': 'int8'}},  # 
  Квантинизация весов    
  {'name': 'quantize_nodes', 'params': {'dtype': 'int8'}},    
  # Квантинизация активаций
]

Перевод переменных в новую битность и сокращает размер нашего файла, а значит и оптимизирует производительность, хотя и слегка понижает точность нейросети. Но эффективность метода в достаточной мере не оценена. Вариант отлично подходит для запуска нейронки на мобильных устройствах, где не хватает вычислительных мощностей.

Преимущества оптимизации очевидны: уменьшение размера модели, улучшение производительности, ускорение обучения и эффективное использование ресурсов, вашего TPU и GPU. Пользуемся и кодим)