Поскольку TF2 ориентирован на активное выполнение, концепция сеанса ушла, и все, что выполняется, превратилось в конкретные функции. Но концепция традиционного замороженного графа все еще может найти свои варианты использования, особенно для разработки и отладки.
Допустим, у вас есть очень простая модель, построенная на основе Keras.
keras = tf.keras
class MyCustomLayer(keras.layers.Layer):
def __init__(self):
super(MyCustomLayer, self).__init__(self)
self._weight = tf.Variable(initial_value=(2., 3.))
def call(self, input):
output = tf.sigmoid(input) * self._weight
return output
model = keras.models.Sequential(
[keras.layers.Input((1,2)), MyCustomLayer()])
В TF1.x вы можете получить сеанс, а затем график через
tf.keras.backend.get_session()
В TF2 вы не можете получить глобальный сеанс, но у вас есть доступ к конкретной функции.
model = tf.keras.Sequential(...)
func = tf.function(model).get_concrete_function(
tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
С конкретной функцией вы уже можете получить определение графика через
func.graph.as_graph_def()
Однако он полон переменных и ReadVariableOp. Нам нужно преобразовать их в постоянные, потому что при выводе модели не нужно обновлять веса.
Это можно сделать с помощью волшебной функции convert_variables_to_constants_v2_as_graph в TF2.2. Раньше он назывался convert_variables_to_constants_v2 в TF2.1 IIRC.
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(func)
Но этого мало. Нам нужно больше оптимизации графа, чтобы получить чистый граф вывода. Давайте вызовем грэпплера.
input_tensors = [ tensor for tensor in frozen_func.inputs if tensor.dtype != tf.resource ] output_tensors = frozen_func.outputs graph_def = run_graph_optimizations( graph_def, input_tensors, output_tensors, config=get_grappler_config(["constfold", "function"]), graph=frozen_func.graph)
Хорошая работа! У вас должен быть чистый график вывода.
Вот полный код:
