TensorFlow pb模型修改和优化

# write pb model
with tf.io.gfile.GFile(model_path, "wb") as f:
    f.write(graph_def.SerializeToString())

# write pbtxt model
tf.io.write_graph(graph_def, os.path.dirname(model_path), os.path.basename(model_path))

import tensorflow as tf
import numpy as np

tf.compat.v1.disable_eager_execution()
tf.compat.v1.reset_default_graph()

m = 200
k = 256
n = 128

a_shape = [m, k]
b_shape = [k, n]

np.random.seed(0)
input_np = np.random.uniform(low=0.0, high=1.0, size=a_shape).astype("float32")
kernel_np = np.random.uniform(low=0.0, high=1.0, size=b_shape).astype("float32")

# 构建模型
pld1 = tf.compat.v1.placeholder(dtype="float32", shape=a_shape, name="input1")
kernel = tf.constant(kernel_np, dtype="float32")
feed_dict = {pld1: input_np}

result_tf = tf.raw_ops.MatMul(a=pld1, b=kernel, transpose_a=False, transpose_b=False)

with tf.compat.v1.Session() as sess:
    results = sess.run(result_tf, feed_dict=feed_dict)
    print("results:", results)

# 保存模型
dump_model_name = "matmul_graph.pb"

graph = tf.compat.v1.get_default_graph()
graph_def = graph.as_graph_def()
with tf.io.gfile.GFile(dump_model_name, "wb") as f:
    f.write(graph_def.SerializeToString())

from google.protobuf import text_format

graph_def = tf.compat.v1.GraphDef()

# read pb model
with tf.io.gfile.GFile(model_path, "rb") as f:
    graph_def.ParseFromString(f.read())

# read pbtxt model
with open(model_path, "r") as pf:
    text_format.Parse(pf.read(), graph_def)

node.name
node.op
node.input
node.device
# please ref https://www.tensorflow.org/api_docs/python/tf/compat/v1/AttrValue
node.attr[attr_name].f # b, i, tensor, etc.

# graph_def中node遍历：
for node in graph_def.node:
    ##

from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import node_def_pb2
from tensorflow.python.framework import tensor_util

pld_node = node_def_pb2.NodeDef()
pld_node.name = name
pld_node.op = "Placeholder"

shape = tf.TensorShape([None, 3, 256, 256])
pld_node.attr["shape"].CopyFrom(attr_value_pb2.AttrValue(shape=shape.as_proto()))
dtype = tf.dtypes.as_dtype("float32")
pld_node.attr["dtype"].CopyFrom(attr_value_pb2.AttrValue(type=dtype.as_datatype_enum))

# other commonly used setting
node.input.extend(in_node_names)
node.attr["value"].CopyFrom(
    attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
        np_array, np_array.type, np_array.shape)))

graph_def_n = tf.compat.v1.GraphDef()
for node in graph_def_o.node:
    node_n = node_def_pb2.NodeDef()
    node_n.CopyFrom(node)
    graph_def_n.node.extend([node_n])

# you probably need copy other value like version, etc. from old graph
graph_def_n.version = graph_def_o.version
graph_def_n.library.CopyFrom(graph_def_o.library)
graph_def_n.versions.CopyFrom(graph_def_o.versions)

return graph_def_n

#!/bin/bash

graphdef=input_model.pb
inputs=Placeholder_1:0,Placeholder_2:0
outputs=output0:0,output1:0

output=${graphdef}.onnx

python -m tf2onnx.convert \
    --graphdef ${graphdef} \
    --output ${output} \
    --inputs ${inputs} \
    --outputs ${outputs}\
    --opset 12