我有一個(TF2)保存模型,其中充滿了訓練操作混亂,我正在嘗試使用 grappler對其進行優化以進行推理,但我想隨后將其保存回 TF2 保存模型(以保持一般作業流程遠離來自 TF1)。
我目前有:
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph
from tensorflow.lite.python.util import run_graph_optimizations, get_grappler_config
# Load the saved-model and get the inference concrete function
sm = tf.saved_model.load('path/to/savedmodel/dir')
func = sm.signatures['serving_default']
# Replace variables with constants in order to get rid of the training clutter
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(func)
# Use grappler to optimize the concrete function graph after replacing vars with constants
input_tensors = [tsr for tsr in frozen_func.inputs if tsr.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)
# Here the intention is to somehow reconvert the optimized graph-def into a concrete function
# and subsequently re-save that as a TF2(not TF1!) saved-model, is there a way to do that?
frozen_func_graph = tf.Graph()
with frozen_func_graph.as_default():
tf.import_graph_def(graph_def, name='')
# ... what now?
問題是,由于TF2中已棄用直接使用 tf.Graph,我打算將優化后的圖轉換回 TF2 保存模型。我正在考慮通過以某種方式手動構建一個包裝這個優化圖的 ConcreteFunction 來做到這一點,但據我研究,現在似乎有辦法做到這一點。這基本上意味著我仍然必須使用 TF1 兼容 API,理想情況下我想避免這種情況。
我真正想避免的丑陋(丑陋)選項是(還沒有嘗試過,但可能會起作用):
- 使用 v1 API 使用tf.compat.v1.saved_model.builder.SavedModelBuilder構造 TF1 保存模型并保存 TF1 保存模型
- 使用 v2 API加載回 TF1 保存模型(因此 tf.saved_model.load 而不是 tf.compat.v1.saved_model.load,前者將 TF1 保存模型自動轉換為 TF2 保存模型)
- (re)保存轉換后的 TF2 保存模型
有沒有辦法很好地做到這一點?如果我不想,最好也不必被迫轉儲優化的保存模型,似乎不可能在記憶體中構建保存模型?(雖然這不是什么大問題)
uj5u.com熱心網友回復:
終于明白了,不理想,因為我使用內部(或多或少)TF2 API 呼叫,但至少根本不使用 TF1 兼容 API。這是完整的代碼。
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph
from tensorflow.lite.python.util import run_graph_optimizations, get_grappler_config
from tensorflow.python.tools.optimize_for_inference_lib import optimize_for_inference
from tensorflow.python.eager import context, wrap_function
# Load the saved-model and get the inference concrete function
sm = tf.saved_model.load('path/to/savedmodel/dir')
func = sm.signatures['serving_default'] # note: key might differ according to what your model's inference function is
# Replace variables with constants in order to get rid of the training clutter
frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(func)
# Use grappler to optimize the concrete function graph after replacing vars with constants
input_tensors = [tsr for tsr in frozen_func.inputs if tsr.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)
# Optimize for inference
input_tsr_names = [tsr.name for tsr in input_tensors]
output_tsr_names = [tsr.name for tsr in output_tensors]
input_node_names = list(set([tsr_name.rsplit(':', 1)[0] for tsr_name in input_tsr_names]))
output_node_names = list(set([tsr_name.rsplit(':', 1)[0] for tsr_name in output_tsr_names]))
graph_def = optimize_for_inference(input_graph_def=graph_def,
input_node_names=input_node_names,
placeholder_type_enum=tf.dtypes.float32.as_datatype_enum,
output_node_names=output_node_names,
toco_compatible=True)
# This next part inspired from _construct_concrete_function function here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/framework/convert_to_constants.py#L1062
# Remove old functions to use updated functions from graph def - not sure if this is actually needed here, didn't look into it
for f in graph_def.library.function:
if context.context().has_function(f.signature.name):
context.context().remove_function(f.signature.name)
# GraphDef to concrete function
opt_frozen_func = wrap_function.function_from_graph_def(graph_def,
input_tsr_names,
output_tsr_names)
# Wrap concrete function into module to export as saved-model
class OptimizedFrozenModel(tf.Module):
def __init__(self, name=None):
super().__init__(name)
module = OptimizedFrozenModel()
module.__call__ = opt_frozen_func
# Export frozen & optimized saved-model
tf.saved_model.save(module, 'path/to/optimized_savedmodel/dir', signatures=opt_frozen_func)
轉載請註明出處,本文鏈接:https://www.uj5u.com/gongcheng/524106.html
