这是我们在 PipelineAI 笔记本中使用的一段代码,用于在 Jupyter 笔记本中内联显示 TensorFlow 图:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import re
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2
def convert_graph_to_dot(input_graph, output_dot, is_input_graph_binary):
graph = graph_pb2.GraphDef()
with open(input_graph, "rb") as fh:
if is_input_graph_binary:
graph.ParseFromString(fh.read())
else:
text_format.Merge(fh.read(), graph)
with open(output_dot, "wt") as fh:
print("digraph graphname {", file=fh)
for node in graph.node:
output_name = node.name
print(" \"" + output_name + "\" [label=\"" + node.op + "\"];", file=fh)
for input_full_name in node.input:
parts = input_full_name.split(":")
input_name = re.sub(r"^\^", "", parts[0])
print(" \"" + input_name + "\" -> \"" + output_name + "\";", file=fh)
print("}", file=fh)
print("Created dot file '%s' for graph '%s'." % (output_dot, input_graph))
input_graph='/root/models/optimize_me/linear/cpu/unoptimized_cpu.pb'
output_dot='/root/notebooks/unoptimized_cpu.dot'
convert_graph_to_dot(input_graph=input_graph, output_dot=output_dot, is_input_graph_binary=True)
使用 graphviz,您可以在笔记本单元格中使用 %%bash 魔法将 .dot 转换为 .png:
%%bash
dot -T png /root/notebooks/unoptimized_cpu.dot \
-o /root/notebooks/unoptimized_cpu.png > /tmp/a.out
最后,在笔记本中显示图表:
from IPython.display import Image
Image('/root/notebooks/unoptimized_cpu.png', width=1024, height=768)
这是一个在 TensorFlow 中实现的简单线性回归模型的示例:
这是用于在生产中部署和服务 TensorFlow 模型的优化版本(也使用上述代码 sn-ps 渲染):
更多此类优化的示例和详细信息,请访问http://pipeline.ai