Замораживание графа в тензорном потоке при использовании tf.image.Dataset

Question

Я использую tensorflow.python.tools.freeze_graph, чтобы заморозить график тензорного потока в функции ниже:

def freeze_and_save_graph(self, session, save_dir, name):
    checkpoint_prefix = os.path.join(save_dir, "model")
    checkpoint_state_name = "checkpoint"
    input_graph_name = "input_graph.pbtxt"
    output_graph_name = name

    # saver = tf.train.Saver(tf.trainable_variables(), max_to_keep=max_checkpoints)
    checkpoint_path = self.saver.save(
        session,
        checkpoint_prefix,
        global_step=0,
        latest_filename=checkpoint_state_name)
    tf.train.write_graph(session.graph, save_dir, input_graph_name, as_text=True)
    input_graph_path = os.path.join(save_dir, input_graph_name)
    input_saver_def_path = ""
    input_binary = False
    output_node_names = "model_1/output"
    restore_op_name = "save/restore_all"
    filename_tensor_name = "save/Const:0"
    output_graph_path = os.path.join(save_dir, output_graph_name)
    clear_devices = False
    freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
                              input_binary, checkpoint_path, output_node_names,
                              restore_op_name, filename_tensor_name,
                              output_graph_path, clear_devices, "")

В последнее время я переключаюсь на использование tensorflow.image.Dataset для предварительной обработки следующим образом:

data = tf.data.Dataset.from_tensor_slices((images_train, onehot_train))
data = data.map(lambda x, y: (preprocessing_fn(x), y), num_parallel_calls=32)
data = data.shuffle(len(images_train))
data = data.batch(batch_size)
data = data.prefetch(5)
iterator = data.make_initializable_iterator()
next_element = iterator.get_next()
init_op = iterator.initializer
session.run(init_op)

После внесения изменений замораживание графика занимает вечность. Размер input_graph.pbtxt увеличился с 500 КБ до 150 МБ. Посмотрев, виновник - два тензора с тем же размером и формой, что и мои тренировочные данные, и с определением tensor_content. То есть данные тренировки были сохранены в файле.

Как сохранить график без этих данных?

Answer 1

Я нашел решение. Используйте заполнители вместо создания набора данных непосредственно из данных. Изменения:

image_tensor = tf.placeholder(tf.float32, shape=self.x_image.shape)
onehot_tensor = tf.placeholder(tf.float32, shape=self.y_true.shape)
data = tf.data.Dataset.from_tensor_slices((image_tensor, onehot_tensor))

и

session.run(init_op, feed_dict={images_tensor: image_train, onehot_tensor: onehot_train})

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