tensorflow estimator的使用

点击下方图片查看HappyChart专业绘图软件



由于以前都是使用tensorflow底层的API，在模型移植方面有一些问题，就考虑到了使用它的高级API estimator。estimator是tensorflow的一个高级API，它的好处就是不用关注底层的物理设备，在cpu，tpu、gpu下运行或者使用分布式训练都不用修改代码。本文档旨在走通estimator的整个流程，其中一些数据处理和图的构建可能不太合理，可根据自己实际情况进行改进。

tensorflow：1.14.0

python: 3.6.5

git代码：tfestimator

构建input_fn

import tensorflow as tf

# 设置log级别
tf.compat.v1.logging.set_verbosity(tf.logging.INFO)

def data_generator():
    with open("./iris.data", 'r', encoding="utf-8") as f:
        for line in f:
            arr = line.strip().split(",")
            yield [arr[:-1]], 1 if arr[-1]=="Iris-setosa" else 2 if arr[-1]=="Iris-versicolor" else 0


def input_fn():
    """处理输入数据，返回一个元组，第一个元素为features，第二个元素为label"""
    output_shape = ([None, None], ())
    output_type = (tf.float32, tf.int32)
    # 通过生成器可以处理大文件
    dataset = tf.data.Dataset.from_generator(data_generator, output_type, output_shape)
    dataset = dataset.repeat(10).shuffle(50).batch(10)
    dataset = dataset.map(map_func=lambda x, y: (x, tf.one_hot(y, 3)))

    return dataset

构建模型model_fn

def model_fn(features, labels, mode, params):
    """
    必须制定四个参数features, labels, mode, params， features和labels是由输入函数input_fn返回的数据， mode为tf.estimator.ModeKeys.TRAIN、tf.estimator.ModeKeys.EVAL和tf.estimator.ModeKeys.PREDICT其中的一个，param为传递的超参数，是一个dict，需要的自己设置。 这里需要注意的一点就是labels的使用要写到ModeKeys.TRAIN或ModeKeys.EVAL的条件内，否则在预测的时候可能会出现问题。
    """
    if isinstance(features, dict):  # 在serving阶段使用
        input = features["input"]
    else:
        input = features
    input = tf.reshape(input, [-1, 4])

    # 模型定义
    dense_1 = tf.layers.dense(input, 10, activation=tf.nn.relu, kernel_initializer=tf.initializers.random_normal())
    dense_2 = tf.layers.dense(dense_1, 10, activation=tf.nn.relu, kernel_initializer=tf.initializers.random_normal())

    w = tf.get_variable("weight", shape=[10, 3], dtype=tf.float32, initializer=tf.variance_scaling_initializer())

    logits = tf.matmul(dense_2, w)

	prob = tf.nn.softmax(logits)
    prob = tf.reduce_max(prob, axis=1)

    # 结果需要返回tf.estimator.EstimatorSpec对象，具体需要设置的东西参见其代码注释
    if mode == tf.estimator.ModeKeys.TRAIN:
        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels, logits)
        loss = tf.reduce_mean(loss)
        # 设置log hook
        logging_hook = tf.estimator.LoggingTensorHook({"loss": loss, "prob": prob}, every_n_iter=1)
        
        optimizer = tf.train.GradientDescentOptimizer(0.01)
        train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode,  loss=loss, train_op=train_op, training_hooks=[logging_hook])

    if mode == tf.estimator.ModeKeys.EVAL:
        loss = tf.nn.softmax_cross_entropy_with_logits_v2(labels, logits)
        loss = tf.reduce_mean(loss)
        return tf.estimator.EstimatorSpec(mode, loss=loss)

    if mode == tf.estimator.ModeKeys.PREDICT:
        pred = tf.argmax(logits, axis=-1)
        predictions = {"pred": pred, "prob": prob}
        return tf.estimator.EstimatorSpec(mode, predictions=predictions)

数据训练及评估

def build_estimator():
    cfg = tf.estimator.RunConfig(save_checkpoints_secs=2)
    estimator = tf.estimator.Estimator(model_fn, model_dir="./my_model", config=cfg, params={})
	
    # 这里往TrainSpec或者EvalSpec传入的input函数都是无参的，所以，如果input函数中有参数时有两种方法解决，1：使用偏函数functools.partial；2：使用lambda表达式再封装input函数一次
    train_spec = tf.estimator.TrainSpec(input_fn)
    eval_spec = tf.estimator.EvalSpec(input_fn=input_fn)

    v = tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
    print(v)

然后运行build_estimator()训练，并在设置的路径"./my_model"下生成一个checkpoint文件。

预测

可以直接加载checkpoint文件进行数据预测：

def pred():
    cfg = tf.estimator.RunConfig(save_checkpoints_secs=2)
    estimator = tf.estimator.Estimator(model_fn, model_dir="./my_model", config=cfg, params={})
    # predict返回一个生成器
    res = estimator.predict(input_fn)

    for i in res:
        print(i)

另一种方式是导出SavedModel做成服务进行预测：

(1) 导出模型

def serving_input_receiver_fn():
    input = tf.placeholder(dtype=tf.float32, shape=[None, 4], name="input")
    # 传入ServingInputReceiver中的参数需要是一个dict，其中key的名称和model_fn中设置服务的features中key的名称一样
    features = {"input": input}
    receiver_tensors = {"input": input}
    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)


def export_model(model_fn, ckp_dir, config, params, save_path):
    """
    model_fn: 模型函数
    ckp_dir: 训练生成的checkpoint文件
    config: 设置的配置
    params: 超参
    save_path: 模型保存地址
    """
    estimator = tf.estimator.Estimator(model_fn, ckp_dir, config=config, params=params)
    estimator.export_saved_model(save_path, serving_input_receiver_fn)

运行export_model(model_fn, "./my_model", None, None, "saved_model")会在saved_model路径下生成一个时间戳文件，文件下包含一个variable文件夹和一个模型saved_model.pb文件，该模型不依赖于具体的语言，可以使用python, java或C++进行预测，或者用于tensorflow-serving。

(2) 进行预测

from pathlib import Path

def predict(fea):
    export_dir = "./saved_model"
    subdirs = [x for x in Path(export_dir).iterdir()
               if x.is_dir() and 'temp' not in str(x)]
    latest = str(sorted(subdirs)[-1])
    print(latest)

    predict_fn = from_saved_model(latest)

    preds = predict_fn({"input": fea})
    print(preds)
    return preds

数据：fea = [[5.0, 2.0, 3.5, 1.0],[6.3, 2.5, 5.0, 1.9]]，运行predict(fea)返回{'prob': array([0.7947242, 0.8389487], dtype=float32), 'pred': array([2, 0], dtype=int64)}

封装服务

(1) 封装为flask服务

from flask import Flask, jsonify

app = Flask(__name__)

from predictor import predict  # 上面的predict函数


@app.route('/serving/<data>')
def serving(data):
    feature = eval(data)
    preds = predict(feature)
    preds = jsonify(str(preds))
    return preds


if __name__ == '__main__':
    app.run()

(2) 使用tensorflow-serving
简便方法是在docker中拉取tensorflow/serving镜像
运行：

docker run -p 8501:8501 --mount type=bind,source=/home/docs/saved_model,target=/models/saved_model -e MODEL_NAME=saved_model -t tensorflow/serving

其中saved_model为模型的名称，需要前后一致。
预测：
[1] 使用命令行

curl -d '{"instances":[[5.0, 2.0, 3.5, 1.0],[6.3, 2.5, 5.0, 1.9]]}' -X POST http://localhost:8501/v1/models/saved_model:predict

结果为：

{
    "predictions": [
        {
            "pred": 2,
            "prob": 0.794724226
        },
        {
            "pred": 0,
            "prob": 0.838948727
        }
    ]
}

[2] 使用restful API

import json
import requests


if __name__ == '__main__':
    d = [[5.0, 2.0, 3.5, 1.0], [6.3, 2.5, 5, 1.9]]
    data = {"instances": d}
    data = json.dumps(data)
    
    response = requests.post("http://localhost:8501/v1/models/saved_model:predict", data=data)
    x = response.content.decode("utf-8")
    print("状态码为：{}".format(response.status_code))
    print("预测结果为：{}".format(json.loads(x)))

结果为：

状态码为：200
预测结果为：{'predictions': [{'pred': 2, 'prob': 0.794724226}, {'pred': 0, 'prob': 0.838948727}]}

[3] 使用gRPC

import grpc
from tensorflow_serving.apis import prediction_service_pb2_grpc, predict_pb2
from tensorflow.python.platform import flags
from tensorflow.python.framework import tensor_util


FLAGS = flags.FLAGS
flags.DEFINE_string("server", "localhost:8500", "PredictionService host:port")


def main():
    channel = grpc.insecure_channel(FLAGS.server)
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    request = predict_pb2.PredictRequest()
    request.model_spec.name = "saved_model"  # 保存模型的名称
    request.model_spec.signature_name = "serving_default"

    input_data = [[5.0, 2.0, 3.5, 1.0], [6.3, 2.5, 5.0, 1.9]]

    tensor_proto = tensor_util.make_tensor_proto(input_data)
    request.inputs["input"].CopyFrom(tensor_proto)  # `input`为模型输入的key

    # result = stub.Predict(request, 5.0)
    future = stub.Predict.future(request, 5)  # 并发运行
    result = future.result()
    print(result)
    res = result.outputs["pred"].int64_val
    prob = result.outputs["prob"].float_val
    
    print("预测结果：{}".format(list(res)))
    print("预测概率：{}".format(list(prob)))


if __name__ == '__main__':
    main()

结果为：

outputs {
  key: "pred"
  value {
    dtype: DT_INT64
    tensor_shape {
      dim {
        size: 2
      }
    }
    int64_val: 2
    int64_val: 0
  }
}
outputs {
  key: "prob"
  value {
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: 2
      }
    }
    float_val: 0.7947242259979248
    float_val: 0.8389487266540527
  }
}
model_spec {
  name: "saved_model"
  version {
    value: 1602828405
  }
  signature_name: "serving_default"
}

预测结果：[2, 0]
预测结果：[0.7947242259979248, 0.8389487266540527]

预测函数from_saved_model定义

由于tensorflow1.14.0中没有了contrib包，其中的预测函数from_saved_model无法使用，所以这里就从tensorflow1.12中的contrib包中摘抄出函数from_saved_model

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import six
import logging

from tensorflow.contrib.saved_model.python.saved_model import reader
from tensorflow.python.client import session
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import loader
from tensorflow.python.saved_model import signature_constants

DEFAULT_TAGS = 'serve'

_DEFAULT_INPUT_ALTERNATIVE_FORMAT = 'default_input_alternative:{}'


@six.add_metaclass(abc.ABCMeta)
class Predictor(object):
  """Abstract base class for all predictors."""

  @property
  def graph(self):
    return self._graph

  @property
  def session(self):
    return self._session

  @property
  def feed_tensors(self):
    return self._feed_tensors

  @property
  def fetch_tensors(self):
    return self._fetch_tensors

  def __repr__(self):
    return '{} with feed tensors {} and fetch_tensors {}'.format(
        type(self).__name__, self._feed_tensors, self._fetch_tensors)

  def __call__(self, input_dict):
    """Returns predictions based on `input_dict`.

    Args:
      input_dict: a `dict` mapping strings to numpy arrays. These keys
        must match `self._feed_tensors.keys()`.

    Returns:
      A `dict` mapping strings to numpy arrays. The keys match
      `self.fetch_tensors.keys()`.

    Raises:
      ValueError: `input_dict` does not match `feed_tensors`.
    """
    # TODO(jamieas): make validation optional?
    input_keys = set(input_dict.keys())
    expected_keys = set(self.feed_tensors.keys())
    unexpected_keys = input_keys - expected_keys
    if unexpected_keys:
      raise ValueError(
          'Got unexpected keys in input_dict: {}\nexpected: {}'.format(
              unexpected_keys, expected_keys))

    feed_dict = {}
    for key in self.feed_tensors.keys():
      value = input_dict.get(key)
      if value is not None:
        feed_dict[self.feed_tensors[key]] = value
    return self._session.run(fetches=self.fetch_tensors, feed_dict=feed_dict)


def get_meta_graph_def(saved_model_dir, tags):
  """Gets `MetaGraphDef` from a directory containing a `SavedModel`.

  Returns the `MetaGraphDef` for the given tag-set and SavedModel directory.

  Args:
    saved_model_dir: Directory containing the SavedModel.
    tags: Comma separated list of tags used to identify the correct
      `MetaGraphDef`.

  Raises:
    ValueError: An error when the given tags cannot be found.

  Returns:
    A `MetaGraphDef` corresponding to the given tags.
  """
  saved_model = reader.read_saved_model(saved_model_dir)
  set_of_tags = set([tag.strip() for tag in tags.split(',')])
  for meta_graph_def in saved_model.meta_graphs:
    if set(meta_graph_def.meta_info_def.tags) == set_of_tags:
      return meta_graph_def
  raise ValueError('Could not find MetaGraphDef with tags {}'.format(tags))


def _get_signature_def(signature_def_key, export_dir, tags):
  """Construct a `SignatureDef` proto."""
  signature_def_key = (
      signature_def_key or
      signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY)

  metagraph_def = get_meta_graph_def(export_dir, tags)

  try:
    signature_def = metagraph_def.signature_def[signature_def_key]
  except KeyError as e:
    formatted_key = _DEFAULT_INPUT_ALTERNATIVE_FORMAT.format(
        signature_def_key)
    try:
      signature_def = metagraph_def.signature_def[formatted_key]
    except KeyError:
      raise ValueError(
          'Got signature_def_key "{}". Available signatures are {}. '
          'Original error:\n{}'.format(
              signature_def_key, list(metagraph_def.signature_def), e))
    logging.warning('Could not find signature def "%s". '
                    'Using "%s" instead', signature_def_key, formatted_key)
  return signature_def


def _check_signature_arguments(signature_def_key,
                               signature_def,
                               input_names,
                               output_names):
  """Validates signature arguments for `SavedModelPredictor`."""
  signature_def_key_specified = signature_def_key is not None
  signature_def_specified = signature_def is not None
  input_names_specified = input_names is not None
  output_names_specified = output_names is not None
  if input_names_specified != output_names_specified:
    raise ValueError(
        'input_names and output_names must both be specified or both be '
        'unspecified.'
    )

  if (signature_def_key_specified + signature_def_specified +
      input_names_specified > 1):
    raise ValueError(
        'You must specify at most one of signature_def_key OR signature_def OR'
        '(input_names AND output_names).'
    )


class SavedModelPredictor(Predictor):
  """A `Predictor` constructed from a `SavedModel`."""

  def __init__(self,
               export_dir,
               signature_def_key=None,
               signature_def=None,
               input_names=None,
               output_names=None,
               tags=None,
               graph=None,
               config=None):
    """Initialize a `CoreEstimatorPredictor`.

    Args:
      export_dir: a path to a directory containing a `SavedModel`.
      signature_def_key: Optional string specifying the signature to use. If
        `None`, then `DEFAULT_SERVING_SIGNATURE_DEF_KEY` is used. Only one of
        `signature_def_key` and `signature_def` should be specified.
      signature_def: A `SignatureDef` proto specifying the inputs and outputs
        for prediction. Only one of `signature_def_key` and `signature_def`
        should be specified.
      input_names: A dictionary mapping strings to `Tensor`s in the `SavedModel`
        that represent the input. The keys can be any string of the user's
        choosing.
      output_names: A dictionary mapping strings to `Tensor`s in the
        `SavedModel` that represent the output. The keys can be any string of
        the user's choosing.
      tags: Optional. Comma separated list of tags that will be used to retrieve
        the correct `SignatureDef`. Defaults to `DEFAULT_TAGS`.
      graph: Optional. The Tensorflow `graph` in which prediction should be
        done.
      config: `ConfigProto` proto used to configure the session.
    Raises:
      ValueError: If more than one of signature_def_key OR signature_def OR
        (input_names AND output_names) is specified.
    """
    _check_signature_arguments(
        signature_def_key, signature_def, input_names, output_names)
    tags = tags or DEFAULT_TAGS
    self._graph = graph or ops.Graph()

    with self._graph.as_default():
      self._session = session.Session(config=config)
      loader.load(self._session, tags.split(','), export_dir)

    if input_names is None:
      if signature_def is None:
        signature_def = _get_signature_def(signature_def_key, export_dir, tags)
      input_names = {k: v.name for k, v in signature_def.inputs.items()}
      output_names = {k: v.name for k, v in signature_def.outputs.items()}

    self._feed_tensors = {k: self._graph.get_tensor_by_name(v)
                          for k, v in input_names.items()}
    self._fetch_tensors = {k: self._graph.get_tensor_by_name(v)
                           for k, v in output_names.items()}


def from_saved_model(export_dir,
                     signature_def_key=None,
                     signature_def=None,
                     input_names=None,
                     output_names=None,
                     tags=None,
                     graph=None,
                     config=None):
  """Constructs a `Predictor` from a `SavedModel` on disk.

  Args:
    export_dir: a path to a directory containing a `SavedModel`.
    signature_def_key: Optional string specifying the signature to use. If
      `None`, then `DEFAULT_SERVING_SIGNATURE_DEF_KEY` is used. Only one of
    `signature_def_key` and `signature_def`
    signature_def: A `SignatureDef` proto specifying the inputs and outputs
      for prediction. Only one of `signature_def_key` and `signature_def`
      should be specified.
      input_names: A dictionary mapping strings to `Tensor`s in the `SavedModel`
        that represent the input. The keys can be any string of the user's
        choosing.
      output_names: A dictionary mapping strings to `Tensor`s in the
        `SavedModel` that represent the output. The keys can be any string of
        the user's choosing.
    tags: Optional. Tags that will be used to retrieve the correct
      `SignatureDef`. Defaults to `DEFAULT_TAGS`.
    graph: Optional. The Tensorflow `graph` in which prediction should be
      done.
    config: `ConfigProto` proto used to configure the session.

  Returns:
    An initialized `Predictor`.

  Raises:
    ValueError: More than one of `signature_def_key` and `signature_def` is
      specified.
  """
  return SavedModelPredictor(
      export_dir,
      signature_def_key=signature_def_key,
      signature_def=signature_def,
      input_names=input_names,
      output_names=output_names,
      tags=tags,
      graph=graph,
      config=config)