Сохранение tf.trainable_variables() с помощью convert_variables_to_constants

У меня есть модель Keras, которую я бы хотел преобразовать в protobuf Tensorflow (например, saved_model.pb).

Эта модель исходит из передачи обучения в сети vgg-19, в которой и головка была отключена и обучена с полностью подключенными + softmax слоями, а остальная часть сети vgg-19 была заморожена.

Я могу загрузить модель в Keras, а затем использовать keras.backend.get_session() для запуска модели в тензорном потоке, генерируя правильные прогнозы:

frame = preprocess(cv2.imread("path/to/img.jpg")
keras_model = keras.models.load_model("path/to/keras/model.h5")

keras_prediction = keras_model.predict(frame)

print(keras_prediction)

with keras.backend.get_session() as sess:

    tvars = tf.trainable_variables()

    output = sess.graph.get_tensor_by_name('Softmax:0')
    input_tensor = sess.graph.get_tensor_by_name('input_1:0')

    tf_prediction = sess.run(output, {input_tensor: frame})
    print(tf_prediction) # this matches keras_prediction exactly

Если я не включаю строку tvars = tf.trainable_variables(), то переменная tf_prediction полностью ошибочна и вообще не соответствует выводу из keras_prediction. Фактически все значения в выходе (одиночный массив с 4 значениями вероятности) одинаковы (~ 0.25, все добавляются к 1). Это заставило меня подозревать, что веса для головы только инициализируются до 0, если tf.trainable_variables() не вызывается первым, что было подтверждено после проверки переменных модели. В любом случае вызов tf.trainable_variables() приводит к правильному прогнозированию тензорного потока.

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

with keras.backend.get_session() as sess:
    tvars = tf.trainable_variables()

    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), ['Softmax'])
    graph_io.write_graph(constant_graph, './', 'saved_model.pb', as_text=False)

Я спрашиваю, как я могу сохранить модель Keras в качестве протобуфа Tensorflow с tf.training_variables() неповрежденным?

Большое спасибо!

Ответы

Ответ 1

Итак, ваш подход замораживания переменных на графике (преобразование в константы) должен работать, но не нужен и сложнее других подходов. (подробнее об этом ниже). Если по какой-то причине вы хотите замораживать график (например, экспортировать на мобильное устройство), мне понадобится дополнительная информация, чтобы помочь отладить, поскольку я не уверен, что неявный материал Keras делает за кулисами с вашим графиком. Однако, если вы хотите просто сохранить и загрузить график позже, я могу объяснить, как это сделать (хотя никаких гарантий того, что все, что делает Keras, не будет испорчено..., рад помочь отладить это).

Таким образом, здесь есть два формата. Один из них - это GraphDef, который используется для контрольной точки, поскольку он не содержит метаданных о входах и выходах. Другой - это MetaGraphDef, который содержит метаданные и граф def, метаданные полезны для предсказания и запускают ModelServer (из tensorflow/serve).

В любом случае вам нужно сделать больше, чем просто вызвать graph_io.write_graph, потому что переменные обычно хранятся вне графика.

В обоих случаях используются библиотеки-обертки. tf.train.Saver в основном используется для сохранения и восстановления контрольных точек.

Однако, поскольку вы хотите предсказание, я бы предложил использовать tf.saved_model.builder.SavedModelBuilder для создания двоичного файла SavedModel. Я привел несколько табличек котлов для этого ниже:

from tensorflow.python.saved_model.signature_constants import DEFAULT_SERVING_SIGNATURE_DEF_KEY as DEFAULT_SIG_DEF
builder = tf.saved_model.builder.SavedModelBuilder('./mymodel')
with keras.backend.get_session() as sess:
  output = sess.graph.get_tensor_by_name('Softmax:0')
  input_tensor = sess.graph.get_tensor_by_name('input_1:0')
  sig_def = tf.saved_model.signature_def_utils.predict_signature_def(
    {'input': input_tensor},
    {'output': output}
  )
  builder.add_meta_graph_and_variables(
      sess, tf.saved_model.tag_constants.SERVING,
      signature_def_map={
        DEFAULT_SIG_DEF: sig_def
      }
  )
builder.save()

После запуска этого кода вы должны иметь файл mymodel/saved_model.pb, а также каталог mymodel/variables/ с protobufs, соответствующими значениям переменных.

Затем, чтобы загрузить модель снова, просто используйте tf.saved_model.loader:

# Does Keras give you the ability to start with a fresh graph?
# If not you'll need to do this in a separate program to avoid
# conflicts with the old default graph
with tf.Session(graph=tf.Graph()):
  meta_graph_def = tf.saved_model.loader.load(
      sess, 
      tf.saved_model.tag_constants.SERVING,
      './mymodel'
  )
  # From this point variables and graph structure are restored

  sig_def = meta_graph_def.signature_def[DEFAULT_SIG_DEF]
  print(sess.run(sig_def.outputs['output'], feed_dict={sig_def.inputs['input']: frame}))

Очевидно, что существует более эффективное предсказание, доступное с помощью этого кода через tenorflow/serve или Cloud ML Engine, но это должно работать. Возможно, что Keras делает что-то под капотом, что также будет мешать этому процессу, и если это так, мы хотели бы услышать об этом (и я хотел бы убедиться, что пользователи Keras также могут заморозить графики, поэтому, если вы хотите отправить мне свой код с полным кодом или что-то, может быть, я найду того, кто знает Keras, чтобы помочь мне отлаживать.)

EDIT: здесь вы можете найти пример конца: https://github.com/GoogleCloudPlatform/cloudml-samples/blob/master/census/keras/trainer/model.py#L85