Способы сохранения и загрузки моделей в Keras

Часто спроектированную и обученную модель требуется сохранять для последующего ее использования в прикладных задачах. Либо, еще в процессе обучения создавать контрольные точки на случай возникновения какого-либо сбоя. Могут быть и другие причины, требующие запись данных на носитель и последующего восстановления модели. На этом занятии мы с вами рассмотрим такой базовый функционал пакета Keras.

Для простоты положим, что у нас простая задача классификации изображений цифр последовательной полносвязной сетью:

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist
 
tf.random.set_seed(1)
 
(x_train, y_train), (x_test, y_test) = mnist.load_data()
 
x_train = x_train.reshape(-1, 784) / 255.0
x_test = x_test.reshape(-1, 784) / 255.0
 
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)
 
 
model = keras.Sequential([
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])
 
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['acc'])
model.fit(x_train, y_train, epochs=5)

В самом простом случае для сохранения модели достаточно вызвать метод:

  • model.save() или tf.keras.models.save_model() – для записи модели на носитель;
  • keras.models.load_model() – для загрузки модели по указанному пути.

Например, это можно сделать так:

model.save('16_model')
model_loaded = keras.models.load_model('16_model')

И, затем, прогнать тестовую выборку через загруженную модель:

model_loaded.evaluate(x_test, y_test)

В результате, на диске в рабочем каталоге появится папка с именем «16_model», в которой будут подкаталоги assets и variables, а также файл saved_model.pb. В этом файле хранится архитектура модели и конфигурация обучения (состояние оптимизатора после 5-й эпохи, величины потерь и метрик). Веса хранятся отдельно в каталоге variables.

Вообще на диске хранится следующая информация:

  • архитектура модели;
  • значения весовых коэффициентов модели;
  • информация о начальной настройки метода compile() (вид оптимизатора, функции потерь и метрик);
  • состояние оптимизатора в момент сохранения модели (для возможности продолжения обучения из текущего состояния).

Все это доступно, если модель сохраняется в новом формате пакета Tensorflow – SavedModel. Если же данные сохранить в прежнем формате h5, то в нем сохраняется только архитектура модели, веса и информация метода compile(). Например, текущее состояние оптимизатора в нем будет отсутствовать и продолжить процесс обучения уже не получится.

Конечно, по умолчанию происходит сохранение данных в новый формат SavedModel. Если же вы все-таки хотите использовать старый формат, то достаточно явно указать расширение файла h5 или же установить параметр save_format='h5':

model.save('16_model_2.h5')
model.save('16_model_3', save_format='h5')
model_loaded = keras.models.load_model('16_model_2.h5')

Разумеется, рекомендуется использовать новый формат данных при сохранении моделей.

Методы get_config() и from_config()

Остается вопрос, как более детально происходит сохранение и загрузка архитектуры модели? Что конкретно хранится на диске? Согласно документации:

https://www.tensorflow.org/guide/keras/save_and_serialize

формат SavedModel хранит имя класса, реализацию функцииcall(), потери, веса и значение конфигурации (если в модели реализован метод get_config). Так вот, при отсутствии явного описания конфигурации, реализация функции call() используется для воссоздания работы модели или слоя (если реконструируется отдельный слой). Чтобы все это было понятнее, давайте рассмотрим пример сохранения и загрузки пользовательского класса модели с полносвязными слоями:

class NeuralNetwork(tf.keras.Model):
    def __init__(self, units):
        super().__init__()
        self.units = units
        self.model_layers = [layers.Dense(n, activation='relu') for n in self.units]
 
    def call(self, inputs):
        x = inputs
        for layer in self.model_layers:
            x = layer(x)
        return x

Мы здесь в конструкторе формируем список полносвязных слоев с функцией активации ReLU, а затем, в методе call() пропускаем входной сигнал последовательно через эти слои.

Далее, мы должны воспользоваться этой моделью, чтобы сформировались весовые коэффициенты:

model = NeuralNetwork([128, 10])
 
y = model.predict(tf.expand_dims(x_test[0], axis=0))
print(y)

И, если теперь выполнить сохранение этой модели:

model.save('16_model')

то в файле saved_model.pb будет храниться реализация метода call(), которая и будет определять структуру модели. Тоесть, послееезагрузки:

model_loaded = keras.models.load_model('16_model')

будет воссоздан класс NeuralNetwork вместе с конкретной реализацией метода call() и теми же самыми весовыми коэффициентами. Именно поэтому, при пропускании того же самого входного сигнала, мы получим абсолютно такие же результаты:

y = model_loaded.predict(tf.expand_dims(x_test[0], axis=0))
print(y)

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

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

class NeuralNetworkLinear(tf.keras.Model):
    def __init__(self, units):
        super().__init__()
        self.units = units
        self.model_layers = [layers.Dense(n, activation='linear') for n in self.units]
 
    def call(self, inputs):
        x = inputs
        for layer in self.model_layers:
            x = layer(x)
        return x

А при загрузке указать, чтобы вместо прежнего класса NeuralNetwork был подставлен новый:

model_loaded = keras.models.load_model('16_model', custom_objects={"NeuralNetwork": NeuralNetworkLinear})

Однако, при запуске программы, мы увидим прежний результат, так как через метод call() будут восстановлены исходные полносвязные слоии конструктор класса вызван не будет. Чтобы полноценно использовать новую модель класса NeuralNetworkLinear при загрузке данных, в файл необходимо сохранять конфигурацию модели, а затем, реконструировать ее по этой конфигурации. Для этого в оба класса добавим два специальных метода:

def get_config(self):
        return {'units': self.units}
 
    @classmethod
    def from_config(cls, config):
        return cls(**config)

Первый метод возвращает словарь конфигурации (в данном случае – это свойство units). А второй метод, используя это свойство, вызывает конструктор текущего класса для воссоздания модели.

Как видите, в данном случае, конфигурация модели – это коллекция из числа нейронов каждого слоя. Такой информации вполне достаточно для воссоздания архитектуры модели.

Как только в классах появляются такие методы, то в выходной файл помещается информации о конфигурации модели (словарь), а при загрузке, автоматически вызывается метод from_config() для воссоздания модели. Благодаря этому, после загрузки в полносвязных слоях будут использоваться линейные функции активации класса NeuralNetworkLinear. Вот общий принцип управления структурой модели через методы get_config() и from_config(). На практике рекомендуется всегда определять эти методы при описании собственных моделей.И, разумеется, они определены для стандартных классов моделей:

  • Sequential.from_config(config);
  • Model.from_config(config).

Ну а у самой модели пакета Keras всегда можно вызвать метод:

config = model.get_config()

Аналогичные операции можно выполнять через еще два метода:

  • json_config = model.to_json() – конфигурация архитектуры в формате JSON;
  • new_model = keras.models.model_from_json(json_config) – загрузкаархитектуры модели из JSON-формата.

Конечно, рассмотренныеметоды get_config() / from_config() и to_json() / model_from_json() отвечают исключительно за структуру модели, но не за данные. Остальные параметры сохраняются независимо от архитектуры модели.

Методы get_weights(), set_weights() и save_weights(), load_weights()

Довольно часто требуется извлекать и задавать коэффициенты для отдельных слоев или модели в целом,  а также сохранять их на диск и загружать с диска. Это достаточно просто сделать, с помощью методов:

  • get_weights() – получение коэффициентов модели / слоя;
  • set_weights() – установка коэффициентов модели / слоя;
  • save_weights() – запись коэффициентов модели / слоя на носитель;
  • set_weights() – загрузка коэффициентов модели / слоя с носителя.

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

model = NeuralNetwork([128, 10])
model2 = NeuralNetwork([128, 10])

Затем, пропустить через них входной сигнал (для формирования весовых коэффициентов в слоях):

y = model.predict(tf.expand_dims(x_test[0], axis=0))
print(y)
y = model2.predict(tf.expand_dims(x_test[0], axis=0))
print(y)

А потом скопировать веса, допустим, из первой модели во вторую:

# считываем и записываем веса только после пропускания через модели входного сигнала
# иначе возникнет ошибка из-за отсутствия начальной инициализации весов
weights = model.get_weights()
model2.set_weights(weights)
 
y = model2.predict(tf.expand_dims(x_test[0], axis=0))
print(y)

Как видим, последний вывод для y совпадает по значениям с первой моделью. Это показывает корректность копирования данных.

Похожие действия можно выполнять и на уровне отдельных слоев, так как они поддерживают те же самые методы get_weights() и set_weights().

Вторая пара методов save_weights() и load_weights() позволяет сохранять веса на диск, а затем, считывать их обратно в модель. Делается это очень простопутем вызова этих методов для модели или слоя:

model.save_weights('model_weights')
model2.load_weights('model_weights')

Здесь мы сохраняем веса первой модели в файл с именем model_weights. В результате на диске появятся два файла:

  • model_weights.data-00000-of-00001 – значения весовых коэффициентов;
  • model_weights.index – индексный вспомогательный файл.

Или же можно воспользоваться прежним форматом HDF5 пакета Keras, если указать у файла явно расширение h5 или hdf5:

model.save_weights('model_weights.h5')
model2.load_weights('model_weights.h5')

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

https://www.tensorflow.org/guide/keras/save_and_serialize

Надеюсь, теперь вы представляете себе, как в самых простых ситуациях можно выполнять сохранение данных модели и их последующее восстановление.

Видео по теме