input_fn

# 对 input_fn 进行一层包装是主要是为了传入数据，可直接定义input_fn
def input_fn_build(data, batch_size=64, epochs=10, shuffle=False):
    def input_fn():
        dataset = tf.data.Dataset.from_tensor_slices(data)   # 预先处理好数据再加载， from_generator 中涉及处理数据会导致GPU等待，影响效率
        # dataset = dataset.map(encode, num_parallel_calls=10) # encode用于对数据处理
        dataset = dataset.repeat(epochs)
        if shuffle:
            dataset = dataset.shuffle(10000)
        dataset = dataset.batch(batch_size).prefetch(2)  # 分布式时直接返回dataset
        iterator = dataset.make_one_shot_iterator()
        features, labels = iterator.get_next()
        return features, labels # 用于单机时使用，也可以直接返回dataset
    return input_fn

model_fn

# 若程序中传入一些外部数据导致构建图有问题，可像input_fn一样，再额外加一层包装
# 使用word2vec初始化时，使用Variables来存储(使用keras的Embedding会有问题,不知道是不是哪里操作错了)
def model_fn(features, labels, mode, params):
    input_a = features['input_a']
    input_b = features['input_b']
    outputs = model(input_a, input_b)  # build model

    if mode == tf.estimator.ModeKeys.TRAIN:
        loss = ...
        train_op = ...
        acc, acc_op = tf.metrics.accuracy(labels=labels, predictions=pred)
        tf.summary.scalar('acc', acc)
        hooks = [tf.train.LoggingTensorHook(tensors={'acc': acc_op, 'loss': loss}, every_n_iter=100), 
                tf.train.StepCounterHook()]
        spec = tf.estimator.EstimatorSpec(
            mode=mode, 
            loss=loss, 
            train_op=train_op, 
            training_hooks=hooks, 
            scaffold=scaffold('model.ckpt') if finetune else None
        )

    elif mode == tf.estimator.ModeKeys.EVAL:
        loss = ...
        eval_metric = {
            'acc': tf.metrics.accuracy(labels=label_, predictions=pred), 
            'auc': tf.metrics.auc(labels=label_, predictions=predict[:, -1]), 
            'recall': tf.metrics.recall(labels=label_, predictions=pred), 
            'precision': tf.metrics.precision(labels=label_, predictions=pred)
        }
        spec = tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric)

    else:
        predictions={'prob': prob}
        export_outputs = {'predictions': tf.estimator.export.PredictOutput(predictions)}
        spec = tf.estimator.EstimatorSpec(mode=mode, predictions=predictions, export_outputs=export_outputs)
    return spec

build model

config = tf.estimator.RunConfig(
    save_checkpoints_secs=10*60,
    keep_checkpoint_max=5
)
if tf.test.is_gpu_available():
    conf = tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=False,
        gpu_options=tf.GPUOptions(allow_growth=True))
    config.replace(session_config=conf)
model = tf.estimator.Estimator(model_fn=model_fn, params=params, model_dir=model_dir, config=config)

train_spec = tf.estimator.TrainSpec(input_fn=input_fn_build(...))
eval_spec = tf.estimator.EvalSpec(input_fn=input_fn_build(...), steps=100)
tf.estimator.train_and_evaluate(estimator=model, train_spec=train_spec, eval_spec=eval_spec)
# model.train(input_fn=input_fn_build(...))   # only train
# result = model.evaluate(input_fn=input_fn_build(...))           
# pred = model.predict(input_fn=input_fn_build(...))    

export model

def serving_input_fn():
    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
        'x': tf.placeholder(tf.int32, [None, 4], name='x')
    })()
    return input_fn

model.export_saved_model(model_dir, serving_input_fn)

feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)  # tf.feature_column
serving_input_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
model.export_savedmodel(model_dir, serving_input_receiver_fn)

查看图

saved_model_cli show --dir . --all

predict

from tensorflow.contrib import predictor

model = predictor.from_saved_model(model_dir)  # 路径为pb模型所在目录
prob = model(
    {'x': [[6.4, 3.2, 4.5, 1.5]]})

init_checkpoint

def get_assignment_map_from_checkpoint(tvars, init_checkpoint):
    """Compute the union of the current variables and checkpoint variables."""
    assignment_map = {}
    initialized_variable_names = {}

    name_to_variable = collections.OrderedDict()
    for var in tvars:
        name = var.name
        m = re.match("^(.*):\\d+$", name)
        if m is not None:
            name = m.group(1)
        name_to_variable[name] = var

    init_vars = tf.train.list_variables(init_checkpoint)

    assignment_map = collections.OrderedDict()
    for x in init_vars:
        (name, var) = (x[0], x[1])
        if name not in name_to_variable:
            continue
        assignment_map[name] = name
        initialized_variable_names[name] = 1
        initialized_variable_names[name + ":0"] = 1

    return (assignment_map, initialized_variable_names)

    # 出现 Assignment map with scope 错误时(https://stackoverflow.com/questions/47867748/transfer-learning-with-tf-estimator-estimator-framework)
    # 可尝试返回 return (name_to_variable, initialized_variable_names)


def scaffold(init_ckp='model.ckpt'):
    tvars = tf.trainable_variables()
    (assignment_map, initialized_variable_names) = get_assignment_map_from_checkpoint(tvars, init_ckp)
    init_op = tf.train.init_from_checkpoint(init_ckp, assignment_map)
    return tf.train.Scaffold()