基于TensorFlow模型文件的静态和动态分析

基于TensorFlow模型文件的静态和动态分析

简介

在往期文章中 我们给你推荐一种TensorFlow模型格式 ,我们介绍过TensorFlow众多的模型格式,不同格式在服务端或者移动端场景下有特殊的优化,但无论用什么工具转都一定包含了Graph(计算图)和Variables(模型权重)的信息。有了Graph和Variables,理论上我们就可以重建模型训练时的场景,甚至通过Graph editor等工具去拓展模型功能做更多有趣的事情。

目前TensorFlow提供了SavedModelBuilder和tf.saved_model.loader.load接口,加载后的Graph可以更新后新建的Session中,进而实现检查模型签名、统计Op信息等静态功能。而由于TensorFlow模型签名输入输出是强类型和固定Shape的,因此我们还能在本地模拟模型上线的Inference接口,从而实现基于本地文件的性能Benchmark等动态功能。

本文将以tfmodel项目为例,介绍基于TensorFlow模型文件的静态和动态分析,项目地址 tobegit3hub/tfmodel 。

模型静态分析

首先TensorFlow的SavedModel的模型格式有一定的规范,同一个目录下可以有多个模型版本(Model version),但需要约定使用相同的模型签名,通过文件的目录结构也可以检查模型是否有效,如下图。

当然文件存在不能保证模型可用,例如saved_model.pb和variables.data都是特殊序列化接口的protobuf文件,要保证模型可用需要用TensorFlow API去加载测试。我们验证过TensorFlow的Python API、C++ API和Java API都提供了可用的tf model loader接口,如果你想写代码只需要下面几行就可以实现了。

try:
  session = tf.Session(graph=tf.Graph())
  tf.saved_model.loader.load(session,
                             [tf.saved_model.tag_constants.SERVING],
                             self.model_version_path)
  return True
except Exception as e:
  logging.error("Fail to validate and get error: {}".format(e))
  return False

静态检查除了检查模型有效性,还可以查看模型的模型签名和Op细节。以检查模型签名为例,TensorFlow Serving官方也推出过简单的Python工具来打印Signature信息,saved_model_cli可以指定模型版本的tag和method打印对应信息(嗯,一个TensorFlow模型目录可以有多个模型版本,一个模型版本可以有多个签名,每个签名还有对应的tag)。

当然,我们也可以自己写代码实现定制化的打印格式,所有的模型签名信息都在tf.saved_model.loader.load返回的meta_graph对象中,是一个protobuf序列化的Python对象,从中获取对应的signature_name、method_name、input和output tensors即可,例如用tfmodel工具返回的工具如下。

如果只是打印模型签名,所有的工具看起来都差不多,实际上加载模型后的Session还可以做更多有用的事情,后面动态分析中就会介绍到。

模型动态分析

所谓模型的动态分析,就是让模型可以运行起来,进而分析出模型的预估性能以及显存占用等指标信息。目前Google的TensorFlow Serving tensorflow/serving 和我们的Simple TensorFlow Serving tobegit3hub/simple_tensorflow_serving 都可加载TensorFlow模型提供API服务,但基于Server的Benchmark在一定程度上受限于框架本身实现的性能。实际上看过两者的实现代码,分别是调用了C++和Python的Session.run()接口而已,底层当然还是TensorFlow的MKL或者CUDA库等实现。

既然如此,我们就可以不依赖TensorFlow Serving或者Simple TensorFlow Serving,在模型分析工具中加载模型,并通过Session.run()去做Inference和预估模型的性能,如果有配套工具还可以检查模型预估过程中每个Op的耗时和显存的使用变化情况。

除此之外,如果训练过程中忘记导出TensorBoard event files了,也可以在分析模型时加载完整的Graph,然后导出到TensorBoard中,由于支持Session.run()功能,还可以通过构造不同分布的测试数据来展示更多训练时未被考虑的指标。

我们组同事还在做基于Graph editor的功能,可以将导出的单机版SavedModel模型修改成分布式版本上线,也可以动态修改Op运行的device从而适配离线训练与在线预估不同的GPU资源情况,当然这个与模型分析无关暂时不细讲了。

使用tfmodel

前面介绍了静态和动态分析模型的基本原理,后面介绍的tfmodel就可以实现几乎所有可能做到的功能(主要是TensorFlow Serving命令行工具也没有做的事情)。tfmodel的安装也很简单,可以使用pip install tensorflow-model或者通过docker run tobegit3hub/tfmodel来运行。

首先是检查TensorFlow模型有效性的功能,通过前面介绍的tf.saved_model.loader.load即可做到,通过命令行的检查可以快速判断模型文件是否可以上线,因为加载模型代码与预估服务是一致的,可以保证检查通过的模型文件都符合上线规范。

其次是检查TensorFlow模型的签名和Op信息,这个其实是受到一个Torch项目的启发 Swall0w/torchstat ,虽然不是官方项目但可以展示模型每一层的input shape、output shape、params、memory、flops等信息。目前tfmodel inspect子命令提供了模型签名的可视化和“伪造”请求数据的功能,未来也可以通过TensorFlow API做更多事情。

然后是对TensorFlow模型的本地Benchmark功能,是通过本地起进程运行Session.run()实现的。前面提到的数据构造功能,在MXNet、PyTorch、Scikit-learn等框架下确实比较难实现的,因为无论还是开源框架的模型格式还是ONNX的通用格式,都很少有把模型签名作为模型一部分导出,目前MXNet和ONNX可以通过额外的一个JSON文件提供。

tfmodel benchmark提供了不同batch size下的性能测试报告,这是在本地无网络通信条件下测试的,理论上性能最接近模型本身的运算性能,而且还可以支持GPU。目前测试用占用本地所有资源,因此我们还提供了docker run tobegit3hub/tfmodel容器,可以在容器中限定运算资源来跑Benchmark。

最后是生成TensorBoard event files功能,TensorBoard绝对是深度学习模型结构可视化以及模型效果可视化的神器,如果训练过程中没有保留event files,可以使用tfmodel工具重新加载模型Graph然后生成TensorBoard依赖的event files。

最近,Google还把以前基于Chrome和Timeline文件的Op cost time可视化功能集成到TensorBoard中,也就是说在TensorBoard可视化的Graph中,还可以看到每一个step的每个op的耗时情况,以及op的device还有哪些op未被调用也画得清清楚楚。之前我们就用于定位一个上游依赖多个op但其中一个被feed_dict拦截传值的问题,这部分代码会影响训练性能但在SavedModel的离线分析中加入分析就非常合适了,示例如下。

tf.summary.scalar('test', tf.constant(1))
merged = tf.summary.merge_all()

with tf.Session() as sess:
  writer = tf.summary.FileWriter("./tensorboard", sess.graph)
  run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
  run_metadata = tf.RunMetadata()

  for i in range(10):
    summary, result = sess.run(
        [merged, h],
        feed_dict={a: 10,
                   d: 100},
        options=run_options,
        run_metadata=run_metadata)

    writer.add_run_metadata(run_metadata, 'step%d' % i)
    writer.add_summary(summary, i)

总结

最后总结一下,TensorFlow因为实现了save和load后Session和Graph对象一致的接口,让开发者只基于模型文件也可以做很多静态和动态的事情,让模型文件本身的校验和性能评估有了更好的做法。当然TensorFlow模型也有一定的缺点,基于Op的粒度就不能像Keras、Torch这种逐个layer展示模型信息,而且手动的device placement在多GPU和分布式的预估中增加了不少开发成本。

大家如果有更多模型评估的点子也可以在评论或者Github中提出,当然除了TensorFlow其他框架要支持类似功能也是相当容易的,首先我们还是考虑把MXNet的signature json写到模型文件中吧。

发表评论

电子邮件地址不会被公开。 必填项已用*标注