tensorflow 模型导出总结

# tf1.0
session = tf.Session()
session.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess=session, save_path=args.save_path)  # 读取保存的模型

├── assets/ # 所需的外部文件，例如说初始化的词汇表文件，一般无
├── assets.extra/ # TensorFlow graph 不需要的文件, 例如说给用户知晓的如何使用SavedModel的信息. Tensorflow 不使用这个目录下的文件。
├── saved_model.pb # 保存的是MetaGraph的网络结构
├── variables # 参数权重，包含了所有模型的变量(tf.Variable objects)参数
    ├── variables.data-00000-of-00001
    └── variables.index

"""tf1.0"""
x = tf.placeholder(tf.float32, [None, 784], name="myInput")
y = tf.nn.softmax(tf.matmul(x, W) + b, name="myOutput")
tf.saved_model.simple_save(
                sess,
                export_dir,
                inputs={"myInput": x},
                outputs={"myOutput": y})

def addNameToTensor(someTensor, theName):
    return tf.identity(someTensor, name=theName)

"""estimator"""
def serving_input_fn():
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='segment_ids')
    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
        'label_ids': label_ids,
        'input_ids': input_ids,
        'input_mask': input_mask,
        'segment_ids': segment_ids,
    return input_fn
if do_export:
   estimator._export_to_tpu = False
   estimator.export_saved_model(Flags.export_dir, serving_input_fn)

import tensorflow.python.saved_model
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
builder = saved_model.builder.SavedModelBuilder(export_path)
signature = predict_signature_def(inputs={'myInput': x},
                                  outputs={'myOutput': y})
""" using custom tag instead of: tags=[tag_constants.SERVING] """
builder.add_meta_graph_and_variables(sess=sess,
                                     tags=["myTag"],
                                     signature_def_map={'predict': signature})
builder.save()

def get_saved_model(bert_config, num_labels, use_one_hot_embeddings):
  tf_config = tf.compat.v1.ConfigProto()
  tf_config.gpu_options.allow_growth = True
  model_file = tf.train.latest_checkpoint(FLAGS.output_dir)
  with tf.Graph().as_default(), tf.Session(config=tf_config) as tf_sess:
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='segment_ids')
    loss, per_example_loss, probabilities, predictions = \
          create_model(bert_config, False, input_ids, input_mask, segment_ids, label_ids,
              num_labels, use_one_hot_embeddings)
    saver = tf.train.Saver()
    print("restore;{}".format(model_file))
    saver.restore(tf_sess, model_file)
    tf.saved_model.simple_save(tf_sess,
            FLAGS.output_dir,
            inputs={
              'label_ids': label_ids,
              'input_ids': input_ids,
              'input_mask': input_mask,
              'segment_ids': segment_ids,
            outputs={"probabilities": probabilities})

import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
export_dir = 'inference/pb2saved'
graph_pb = 'inference/robert_tiny_clue/frozen_model.pb'
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
with tf.gfile.GFile(graph_pb, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
sigs = {}
with tf.Session(graph=tf.Graph()) as sess:
    # name="" is important to ensure we don't get spurious prefixing
    tf.import_graph_def(graph_def, name="")
    g = tf.get_default_graph()
    input_ids = sess.graph.get_tensor_by_name(
    "input_ids:0")
    input_mask = sess.graph.get_tensor_by_name(
    "input_mask:0")
    segment_ids = sess.graph.get_tensor_by_name(
    "segment_ids:0")
    probabilities = g.get_tensor_by_name("loss/pred_prob:0")
    sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
        tf.saved_model.signature_def_utils.predict_signature_def(
                "input_ids": input_ids,
                "input_mask": input_mask,
                "segment_ids": segment_ids
            }, {
                "probabilities": probabilities
    builder.add_meta_graph_and_variables(sess,
                                         [tag_constants.SERVING],
                                         signature_def_map=sigs)
builder.save()

model.save('saved_model/my_model')  
"""saved as SavedModel by default"""

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['input_ids'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 128)
        name: input_ids:0
    inputs['input_mask'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 128)
        name: input_mask:0
    inputs['label_ids'] tensor_info:
        dtype: DT_INT32
        shape: (-1)
        name: label_ids:0
    inputs['segment_ids'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 128)
        name: segment_ids:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['probabilities'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 7)
        name: loss/pred_prob:0
  Method name is: tensorflow/serving/predict

predict_fn = tf.contrib.predictor.from_saved_model(args_in_use.model)
prediction = predict_fn({
                "input_ids": [feature.input_ids],
                "input_mask": [feature.input_mask],
                "segment_ids": [feature.segment_ids],
probabilities = prediction["probabilities"]

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ["serve"], export_path)
    graph = tf.get_default_graph()
    feed_dict = {"input_ids_1:0": [feature.input_ids],
                "input_mask_1:0": [feature.input_mask],
                "segment_ids_1:0": [feature.segment_ids]}
    # alternative way
    feed_dict = {sess.graph.get_tensor_by_name("input_ids_1:0"): 
                          [feature.input_ids],
                sess.graph.get_tensor_by_name("input_mask_1:0"):
                          [feature.input_mask],
                sess.graph.get_tensor_by_name("segment_ids_1:0"):
                          [feature.segment_ids]}
    sess.run('loss/pred_prob:0',
               feed_dict=feed_dict

import org.tensorflow.*
SavedModelBundle savedModelBundle = SavedModelBundle.load("./export_path", "serve");
Graph graph = savedModelBundle.graph();
Tensor tensor = this.savedModelBundle.session().runner()
                .feed("input_ids", inputIdTensor)
                .feed("input_mask", inputMaskTensor)
                .feed("segment_ids", inputSegmentTensor)
                .fetch("loss/pred_prob")
                .run().get(0);

$ saved_model_cli show --dir export/1524906774 \
  --tag_set serve --signature_def serving_default
The given SavedModel SignatureDef contains the following input(s):
  inputs['inputs'] tensor_info:
      dtype: DT_STRING
      shape: (-1)
The given SavedModel SignatureDef contains the following output(s):
  outputs['classes'] tensor_info:
      dtype: DT_STRING
      shape: (-1, 3)
  outputs['scores'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 3)
Method name is: tensorflow/serving/classify
$ saved_model_cli run --dir export/1524906774 \
  --tag_set serve --signature_def serving_default \
  --input_examples 'inputs=[{"SepalLength":[5.1],"SepalWidth":[3.3],"PetalLength":[1.7],"PetalWidth":[0.5]}]'
Result for output key classes:
[[b'0' b'1' b'2']]
Result for output key scores:
[[9.9919027e-01 8.0969761e-04 1.2872645e-09]]

"""tf1.0"""
from tensorflow.python.framework.graph_util import convert_variables_to_constants
output_graph_def = convert_variables_to_constants(
                    session,
                    session.graph_def,
                    output_node_names=['loss/pred_prob'])
tf.train.write_graph(output_graph_def, args.export_dir, args.model_name, as_text=False)

"""
NB：首先我们要在create_model() 函数中，为我们需要的输出节点取个名字，
  比如说我们要： probabilities = tf.nn.softmax(logits, axis=-1, name='pred_prob')
def get_frozen_model(bert_config, num_labels, use_one_hot_embeddings):
  tf_config = tf.compat.v1.ConfigProto()
  tf_config.gpu_options.allow_growth = True
  output_node_names = ['loss/pred_prob']
  model_file = tf.train.latest_checkpoint(FLAGS.output_dir)
  with tf.Graph().as_default(), tf.Session(config=tf_config) as tf_sess: 
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='segment_ids')
    create_model(bert_config, False, input_ids, input_mask, segment_ids, label_ids,
            num_labels, use_one_hot_embeddings)
    saver = tf.train.Saver()
    print("restore;{}".format(model_file))
    saver.restore(tf_sess, model_file)
    tmp_g = tf_sess.graph.as_graph_def()
    if FLAGS.use_opt:
      input_tensors = [input_ids, input_mask, segment_ids]
      dtypes = [n.dtype for n in input_tensors]
      print('optimize...')
      tmp_g = optimize_for_inference(tmp_g,
                                    [n.name[:-2] for n in input_tensors],
                                     output_node_names,
                                     [dtype.as_datatype_enum for dtype in dtypes],
                                     False)
    print('freeze...')
    frozen_graph = tf.graph_util.convert_variables_to_constants(tf_sess, 
            tmp_g, output_node_names)
    out_graph_path = os.path.join(FLAGS.output_dir, "frozen_model.pb")
    with tf.io.gfile.GFile(out_graph_path, "wb") as f:
      f.write(frozen_graph.SerializeToString())      
    print(f'pb file saved in {out_graph_path}')

from tensorflow.python.tools import freeze_graph
from tensorflow.python.saved_model import tag_constants
input_saved_model_dir = "./1583934987/"
output_node_names = "loss/pred_prob"
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
output_graph_filename='frozen_graph.pb'
freeze_graph.freeze_graph(input_graph_filename,
  input_saver_def_path,
  input_binary,
  checkpoint_path,
  output_node_names,
  restore_op_name,
  filename_tensor_name,
  output_graph_filename,
  clear_devices,
  "", "", "",
  input_meta_graph,
  input_saved_model_dir,
  saved_model_tags)

from keras import backend as K
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
    Freezes the state of a session into a pruned computation graph.
    Creates a new computation graph where variable nodes are replaced by
    constants taking their current value in the session. The new graph will be
    pruned so subgraphs that are not necessary to compute the requested
    outputs are removed.
    @param session The TensorFlow session to be frozen.
    @param keep_var_names A list of variable names that should not be frozen,
                          or None to freeze all the variables in the graph.
    @param output_names Names of the relevant graph outputs.
    @param clear_devices Remove the device directives from the graph for better portability.
    @return The frozen graph definition.
    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ""
        frozen_graph = tf.graph_util.convert_variables_to_constants(
            session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph
frozen_graph = freeze_session(K.get_session(),
                              output_names=[out.op.name for out in model.outputs])
tf.train.write_graph(frozen_graph, "some_directory", "my_model.pb", as_text=False)

freeze_graph --input_checkpoint model.ckpt-16329 \
             --output_graph 0316_roberta.pb \
             --output_node_names loss/pred_prob \
             --checkpoint_version 1 \
             --input_meta_graph model.ckpt-16329.meta \
             --input_binary true

import tensorflow as tf
def printTensors(pb_file):
    """read pb into graph_def"""
    with tf.gfile.GFile(pb_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    """import graph_def"""
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def)
    """print operations"""
    for op in graph.get_operations():
        print(op.name)
printTensors("path-to-my-pbfile.pb")

import/input_ids:0
import/input_mask:0
import/segment_ids:0
import/loss/pred_prob:0

def addNameToTensor(someTensor, theName):
    return tf.identity(someTensor, name=theName)

with tf.Graph().as_default():
    output_graph_def = tf.GraphDef()
   load pb model
    with open(args_in_use.model, 'rb') as f:
        output_graph_def.ParseFromString(f.read())
        tf.import_graph_def(output_graph_def, name='') #name是必须的
    enter a text and predict
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        input_ids = sess.graph.get_tensor_by_name(
            "input_ids:0")
        input_mask = sess.graph.get_tensor_by_name(
            "input_mask:0")
        segment_ids = sess.graph.get_tensor_by_name(
            "segment_ids:0")
        output = "loss/pred_prob:0"
        feed_dict = {
            input_ids: [feature.input_ids],
            input_mask: [feature.input_mask],
            segment_ids: [feature.segment_ids],
        # 也可以直接使用
        # feed_dict = {
        #     "input_ids:0": [feature.input_ids],
        #     "input_mask:0": [feature.input_mask],
        #     "segment_ids:0": [feature.segment_ids],
        y_pred_cls = sess.run(output, feed_dict=feed_dict)

// TensorUtil.class
public static Session generateSession(String modelPath) throws IOException {
    Preconditions.checkNotNull(modelPath);
    byte[] graphDef = ByteStreams.toByteArray(TensorUtil.class.getResourceAsStream(modelPath));
    LOGGER.info("Graph Def Length: {}", graphDef.length);
    Graph graph = new Graph();
    graph.importGraphDef(graphDef);
    return new Session(graph);
// model.class
this.session = TensorUtil.generateSession(modelPath);
Tensor tensor = this.session.runner()
                    .feed("input_ids", inputIdTensor)
                    .feed("input_mask", inputMaskTensor)
                    .feed("segment_ids", inputSegmentTensor)
                    .fetch("loss/pred_prob")
                    .run().get(0);

"""默认1.0 是HDF5，但是2.0中，是SavedModel，所以需要显性地指定`.h5`后缀"""
model.save('my_model.h5') 

"""keras 1.0"""
model.save_weights('my_model_weights.h5')

"""keras 1.0"""
from keras.models import load_model
model = load_model(model_path)

"""keras 2.0"""
new_model = tf.keras.models.load_model('my_model.h5')

dependencies = {'MyLayer': MyLayer(), 'auc': auc, 'log_loss': log_loss}
model = load_model(model_path, custom_objects=dependencies, compile=False)

"""
To save custom objects to HDF5, you must do the following:
1. Define a get_config method in your object, and optionally a from_config classmethod.
get_config(self) returns a JSON-serializable dictionary of parameters needed to recreate the object.
from_config(cls, config) uses the returned config from get_config to create a new object. By default, this function will use the config as initialization kwargs (return cls(**config)).
2. Pass the object to the custom_objects argument when loading the model. The argument must be a dictionary mapping the string class name to the Python class. E.g. tf.keras.models.load_model(path, custom_objects={'CustomLayer': CustomLayer})

model.load_weights('my_model_weights.h5')

model.load_weights('my_model_weights.h5', by_name=True)

"""
--saved_model_dir:  Type: string. Specifies the full path to the directory containing the SavedModel generated in 1.X or 2.X.
--output_file: Type: string. Specifies the full path of the output file.
tflite_convert \
    --saved_model_dir=1583934987 \
    --output_file=rbt.tflite

tflite_convert --graph_def_file albert_tiny_zh.pb \
               --input_arrays 'input_ids,input_masks,segment_ids' \
               --output_arrays 'finetune_mrc/add, finetune_mrc/add_1'\
               --input_shapes 1,512:1,512:1,512 \
               --output_file saved_model.tflite \
               --enable_v1_converter \
               --experimental_new_converter