如何解决TensorFlow 2:重新保存SavedModel吗?
我正在尝试加载以SavedModel格式保存的模型,然后在其上应用一些计算并重新保存整个管道。最少的代码如下(来自this Kaggle kernel):
# Load the model. After that,the `delg_model` will be of the type
# `tensorflow.python.training.tracking.tracking.AutoTrackable`
delg_model = tf.saved_model.load('path/to/saved/model/dir')
# we don't need the whole model,so we prune it. After that,the
# `global_feature_extraction_fn` will be of the type
# `tensorflow.python.eager.wrap_function.WrappedFunction`
delg_input_tensor_names = ['input_image:0','input_scales:0']
global_feature_extraction_fn = delg_model.prune(
delg_input_tensor_names,['global_descriptors:0'])
问题:现在,我想保存
global_feature_extraction_fn
,并使用其他TF操作以相同的SavedModel格式对输出进行后处理。正确的方法是什么?
我尝试过的
我尝试遵循saving custom models to SavedModel format的TensorFlow文档并定义tf.Module
:
class DelgModule(tf.Module):
def __init__(self):
super().__init__()
@tf.function(input_signature=[
tf.TensorSpec(shape=[None,None,3],name='input_image')
])
def call(self,input_tensor):
# custom function on top of the model's output
embedding = tf.nn.l2_normalize(
global_feature_extraction_fn(
input_tensor,# input_image
tf.convert_to_tensor([0.7,1.0,1.4]) # input_scales
)[0],axis=1,name='l2_normalized_output')
return output_tensors = {
'global_descriptor': embedding
}
delg_module = DelgModule()
然后我在测试图像上运行它以构建tf.function
并确保其正确运行(产生正确的输出)。但是当我尝试将其保存如下时:
tf.saved_model.save(
delg_module,export_dir='./delg_resaved',signatures={
'serving_default': delg_module.call
})
结果模型不正确。原始的一个权重为90 MB,而delg_resaved
中的一个权重仅为800 KB。我还尝试在tf.saved_model.load
函数内进行DelgModule.call
内部,以便图形创建和变量加载完全在该tf.function
内部完成,但是结果保持不变
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。