Хитрости в Machine Learning — работа с несколькими моделями в Keras

Это моя первая статья по теме Машинное обучение. С недавнего времени я профессионально занимаюсь машинным обучением и компьютерным зрением. В этой и будущих статьях я буду делиться наблюдениями и решениями специфических проблем при использовании TensorFlow и Keras. В этой статье я расскажу об одном неочевидном вопросе при работе с TensorFlow и Keras — одновременная загрузка и выполнение нескольких моделей. Если вы не знакомы с тем как работают TensorFlow и Keras внутри, эта тема может стать проблемой для начинающих. Если вас заинтересовала тема, прошу под кат.

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

Подробнее о графах и сессиях можно прочитать на Медиуме и Хабре: здесь и здесь.

Обычно мы работаем с одной моделью и проблем здесь не возникает. Теперь давайте представим, что мы работаем с двумя классами. Оба класса работают с моделями Keras: создают архитектуру нейронки, загружают обученные веса, выполняют предсказание. При выполнении двух классов в одном пайплайне (например, на первом шаге мы выполняем детекцию лица на фото, на следующем — распознавание человека) может возникнуть подобная ошибка:

Error Tensor("norm_layer/l2_normalize:0", shape=(?, 128), dtype=float32) is not an element of this graph

Причина ошибки в том, что Keras по умолчанию работает только с сессией по умолчанию (default session) и не регистрирует новую сессию как сессию по умолчанию.
При работе с моделью Keras пользователь должен явно назначить новую сессию как сессию по умолчанию. Это можно сделать вот так:

self.graph = tf.Graph()
with self.graph.as_default():
    self.session = tf.Session(graph=self.graph)
    with self.session.as_default():
        self.model = WideResNet(face_size, depth=depth, k=width)()
        model_dir = <model_path>
        ...
        self.model.load_weights(fpath)

Мы создаем новый TensorFlow Graph и Session и загружаем модель внутри новой сессии TensorFlow.
Строка

with self.graph.as_default():

означает, что мы хотим использовать новый граф () как граф по умолчанию и в строке

with self.session.as_default():

мы указываем, что хотим использовать сессию self.session как сессию по умолчанию и выполнять последующий код в рамках этой сессии. Конструкция with создает менеджер контекста, который позволяет эффективно работать с памятью, когда мы имеем дело с ресурсоемкими объектами (например, чтение файлов), поскольку автоматически освобождает ресурсы по выходу из блока with.

Когда нам будет нужно выполнить предсказание делаем это так:

with self.graph.as_default():
    with self.session.as_default():
        result = self.model.predict(np.expand_dims(img, axis=0), batch_size=1)

Просто вызываем метод predict() внутри сессии TF, созданной ранее.

Пока на этом все. Всем удачи и до новых встреч!