tensorflow contrib_Tensorflow笔记：你都会了么？

前言

个人在日常工作学习中积累了一些关于Tensorflow的经验，我发现网上的资料坑多又乱，而且不成体系不够全面，官方文档又多又杂，而且介绍点到为止，很多细节上的东西根本不讲。所以整理了这篇笔记，建议点赞收藏（点赞的可以把封面娶走）！（由于篇幅限制，部分代码只能省略，想看完整代码可以看我的其他文章）读了这篇Tensorflow笔记你能够：

1. 成体系的了解Tensorflow的用法（低阶API+Estimator+keras，从训练到部署）。
2. 对于一些博客/文档中没有介绍的重要细节有所了解。

目录：

1. 低阶API
   1. 基础用法
   2. TFrecord数据的制作与读取
   3. 模型加载、保存与Fine-tune
   4. 通过tf.Serving+Docker部署
   5. 分布式训练
2. 高级封装——Estimator
   1. input_fn
   2. model_fn
   3. main
   4. 分布式训练
3. 高级封装——keras
   1. 贯序结构
   2. 复杂结构
   3. 保存与加载
   4. 迁移学习
   5. 部署
   6. 并行训练

一、低阶API

1.1 基础用法

1.1.1 张量

张量（tensor）就是多维向量，维数不定。在Tensorflow中，tensor是最基本的计算单元。

import tensorflow as tf
# 声明一个常量tensor
a = tf.constant(name="a", shape=[2, 2], [[0.0, 1.0], [1.0, 2.0]], dtype=tf.float64)
# 声明一个变量tensor
b = tf.get_variable(name="b", shape=[2, 2],initializer=tf.glorot_normal_initializer(), dtype=tf.float32)
# 通过计算得到tensor
c = tf.add(a, b, name="add")    # 报错，因为a和b的dtype不同，虽然都是float但是仍然不能运算

既然是多维变量，那上面的shape和dtype很容易理解，initializer是对于变量而言的，指依靠哪种分布的生成器来生成变量。

而name就比较复杂了，按理说我们已经有了“变量名”了，一般是不需要在指定一个tensor的name了，可是为什么这里还要又一个name参数呢？但是因为在Tensorflow中有时会有一些tensor只能通过name获取。比如在保存模型和加载模型时，当我们加载别人训练好的模型时，我们可能没有源码，这样就无法通过作者在写代码的时候用的变量名来获取变量。这时我们就要用到这个name了。

另一个不得不用name的情况，就是在写代码的时候，我们将其他的tensor赋值给了已经用过的变量名。下面就是一个例子：

# 初始化一个常数tensor
x = tf.constant(name="x_origin", shape=[1], value=[1.0], dtype=tf.float32) 
# 不断自增
x = tf.add(x, x, name="x_1")
x = tf.add(x, x, name="x_2")
x = tf.add(x, x, name="x_3")

这时如果我们只通过变量名“x”来获取tensor，只能获取name="x_3"的tensor，如果想看中间过程的话，就必须依赖name了，通过get_tensor_by_name函数来获取tensor：

x_2 = tf.get_default_graph().get_tensor_by_name("x_2:0")
print(x_2)

实际应用中，假设一个网络深度达到50层，那可能将一个input张量经过50次全链接，这种情况下如果给每一层都赋一个变量名，就不得不这样写

# 不用 name，只通过变量名区分tensor
x_deep = x_input
x_deep_1 = tf.contrib.layers.fully_connected(inputs=x_deep, num_outputs=128)
"""  ... （这里省略48行） ...  """
x_deep_50 = tf.contrib.layers.fully_connected(inputs=x_deep_49, num_outputs=128)

但是如果有了name来区分tensor，就可以简化成这样

# 通过name来区分tensor
x_deep = x_input
for i in range(50):
    x_deep = tf.contrib.layers.fully_connected(inputs=x_deep, num_outputs=128, name="deep_%d" % i)

所以，我们在刚接触Tensorflow的时候，千万不要忽略张量name的重要性，养成为每一个tensor赋name的习惯。

1.1.2 计算图

一个计算图的例子

前面介绍了张量，每一个张量就对应了一个节点。比如图中x和y就是两个张量，绿色节点是将x和y经过“+”操作后得到的张量，同理橙色节点就是z和(x+y)两个张量经过“x”操作后得到的张量。

提到计算图，就不得不提到一个概念——惰性计算。所谓惰性计算就是当我们定义好一个计算图的时候，每一个张量中是没有数据的，只是储存了“计算方式”，当我们需要得到某一个计算节点（张量）的结果时，Tensorflow会从这个节点向前追溯。比如我们需要计算绿色节点的结果，Tensorflow首先会看需要什么输入，发现时x和y，然后去找x和y节点需要什么输入，到最后到输入节点，就需要人为feed给他了。这个过程中并没有用到的图结构不参与计算（比如z张量以及后面的橙色节点），避免了无效计算。

with tf.Session() as sess:
    # 指定绿色节点的张量作为输出。因为没用到z，所以不需要feed z张量。
    sess.run(绿色节点, feed_dict={x:数据, y:数据})

1.1.3 命名空间

通过命名空间，就可以让张量的命名看起来更好看。

# 无scope（默认scope）声明tensor
weight_layer1 = tf.get_variable(name="weight_layer1", shape=[2, 2],initializer=tf.glorot_normal_initializer())
# 无scope（默认scope）获取tensor
tf.get_default_graph().get_tensor_by_name("weight_layer1:0")

# 有scope声明tensor
with tf.variable_scope("layer1"):
    weight = tf.get_variable(name="weight", shape=[2, 2],initializer=tf.glorot_normal_initializer())
# 有scope获取tensor
tf.get_default_graph().get_tensor_by_name("layer1/weight:0")

1.1.4 会话

人为对计算图feed数据的环境，就是会话。所有跟数据有关的操作，比如训练、保存带变量的网络、加载带变量的网络等等都需要在会话环境中进行。

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(xxxx, feed_dict=(xxx:xxx, xxx:xxx))

1.1.5 样例

这一节我们来看一个整体的例子，用DNN网络对MNIST数据集进行分类：

# 获取数据
(train_images, train_labels), (valid_images, valid_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(-1, 28, 28, 1)
valid_images = valid_images.reshape(-1, 28, 28, 1)
train_images, valid_images = train_images / 255.0, valid_images / 255.0

# 搭建网络结构
input_image = tf.placeholder(name="input_image", shape=[None, 28, 28], dtype=tf.float32)
input_dense = tf.reshape(input_image, shape=[None, 28*28])
input_label = tf.placeholder(name="input_label", shape=[None, 1], dtype=tf.float32)

# 第一层（隐藏层）
with tf.variable_scope("layer1"):
    weights = tf.get_variable(name="weights", shape=[28*28, 128], initializer=tf.glorot_normal_initializer())
    biases = tf.get_variable(name="biases", shape=[128], initializer=tf.glorot_normal_initializer())
    layer1 = tf.nn.relu(tf.matmul(input_dense, weights) + biases, name="layer1")
# 第二层（输出层）
with tf.variable_scope("layer2"):
    weights = tf.get_variable(name="weights", shape=[128, 10], initializer=tf.glorot_normal_initializer())
    biases = tf.get_variable(name="biases", shape=[10], initializer=tf.glorot_normal_initializer())
    y = tf.add(tf.matmul(layer1, weights), biases, name="y")
pred = tf.argmax(y, axis=1, name="pred")

# 损失和优化
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=input_label)
loss = tf.reduce_mean(cross_entropy)
train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss, global_step=tf.train.get_global_step())

# 会话
batch_size = 32
with tf.Session() as sess:
    for i in range(1000):
        x_feed = train_images[i*batch_size, (i+1)*batch_size-1]
        y_feed = train_labels[i*batch_size, (i+1)*batch_size-1]
        _, loss_value = sess.run([train_op, loss], feed_dict={input_image:x_feed, input_label:y_feed})
        if i%1000 == 0:
            saver.save(sess, "./model_ckpt/", global_step=tf.train.get_global_step())

1.2 TFrecord数据的制作与读取

1.2.1 制作

以MNIST数据集为例（不论文本、图片、声音，都是先转化成numpy，在转化成TFRecord），在这里下载好之后，还需要像这样预处理一下。下一步就是把每一张图片读成numpy再写入TFRecord了。读成numpy的过程因人而异因项目而异，个人比较喜欢通过手动制作一个索引文件来读取。具体说来就是用一个文本文件，每行存放一个样本的label、图片路径等信息。大概长这样：

label,file
5,~/data/Mnist/0.png
0,~/data/Mnist/1.png
4,~/data/Mnist/2.png
1,~/data/Mnist/3.png
... ...

这样做的好处是，可以不用一口气把数据读进内存，对于大数据集任务比较友好。而且在多模态的任务中，通过“索引文件”的方式也能够使多种形式的多个文件的读取更加简洁，灵活。

import numpy as np
from PIL import image
import tensorflow as tf

index_file = "./index_file.csv"
writer = tf.python_io.TFRecordWriter("./data/mnist.tfrecord")    # 打开文件

index_list = open(index_file, "r").readlines()[1:]    # 读取索引文件，去掉首行
for line in index_list:
    # 获取label和图片的numpy形式
    label = int(line.split(",")[0])    # 将每行第一个元素读成int，作为label
    img = np.array(Image.open(line.split(",")[1]))    # 根据每行中文件名读取文件，并转化为numpy
    
    # 将label和img捏在一起
    example = tf.train.Example(features=tf.train.Features(feature={
        "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[int(label)])),
        "image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[img.tobytes()])),
    }))  # example对象对label和img数据进行封装

    # 将构建好的 example 写入到 TFRecord
    writer.write(example.SerializeToString())
# 关闭文件
writer.close()

这个过程很简单，先把数据转化成numpy，但后在转化为example，最后写入文件。但是有一个地方需要说明一下，构建example的时候，这个tf.train.Feature()函数可以接收三种数据：

• bytes_list: 可以存储string 和byte两种数据类型。
• float_list: 可以存储float(float32)与double(float64) 两种数据类型。
• int64_list: 可以存储：bool, enum, int32, uint32, int64, uint64。

对于只有一个值（比如label）可以用float_list或int64_list，而像图片、视频、embedding这种列表型的数据，通常转化为bytes格式储存。

1.2.2 读取

TFRecord做好了，要怎么读取呢？我们可以通过tf.data来生成一个迭代器，每次调用都返回一个大小为batch_size的batch。

def read_and_decode(filenames, batch_size=32, num_epochs=None, perform_shuffle=False):
    def _parse_fn(record):
        features = {
            "label": tf.FixedLenFeature([], tf.int64),
            "image": tf.FixedLenFeature([], tf.string),
        }
        parsed = tf.parse_single_example(record, features)
        # image
        image = tf.decode_raw(parsed["image"], tf.uint8)
        image = tf.reshape(image, [28, 28])
        # label
        label = tf.cast(parsed["label"], tf.int64)
        return {"image": image}, label

    # Extract lines from input files using the Dataset API, can pass one filename or filename list
    dataset = tf.data.TFRecordDataset(filenames).map(_parse_fn, num_parallel_calls=10).prefetch(500000)    # multi-thread pre-process then prefetch

    # Randomizes input using a window of 256 elements (read into memory)
    if perform_shuffle:
        dataset = dataset.shuffle(buffer_size=256)

    # epochs from blending together.
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size) # Batch size to use

    iterator = dataset.make_one_shot_iterator()
    batch_features, batch_labels = iterator.get_next()
    return batch_features, batch_labels

对于不同的数据，只需要改动_parse_fn函数就可以。这里有一点很重要，就是在_parse_fn函数中，tf.decode_raw的第二个参数（解码格式），必须和保存TFRecord时候的numpy的格式是一样的，否则会报TypeError，我们保存图片时候采用的是np.uint8，这里解码的时候也要用tf.uint8。

batch_features, batch_labels = read_and_decode("./data/mnist.tfrecord")
with tf.Session() as sess:
    print(sess.run(batch_features["image"][0]))
    print(sess.run(batch_labels[0]))

1.2.3 使用

会写会读之后，我们来简单尝试下怎么用吧！假设我们要用简单的DNN预测MNIST的label。

# 调用 read_and_decode 获取一个 batch 的数据
batch_features, batch_labels = read_and_decode("./data/mnist.tfrecord")

# input
X = tf.cast(batch_features["image"], tf.float32, name="input_image")
X = tf.reshape(X, [-1, 28*28]) / 255    # 将像素点的值标准化到[0,1]之间
label = tf.one_hot(tf.cast(batch_labels, tf.int32, name="input_label"), depth=10, name="label")

# 省略网络结构，直接输出y
y = ... 
pred = tf.nn.softmax(y, name="pred")
# 构建损失
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=label))
# 构建train_op
train_op = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8).minimize(loss)

上面就是简单的，通过read_and_decode函数读取数据，每次sess.run()时，根据向前追溯的计算逻辑，都会自动的调用一次read_and_decode获得一个batch的数据，所以就不需要手动feed数据。

1.3 模型加载、保存与Fine-tune

1.3.1 保存

Tensorflow的保存分为三种：1. checkpoint模式；2. pb模式；3. saved_model模式。

首先假定我们已经有了这样一个简单的线性回归网络结构：

import tensorflow as tf
size = 10
# 构建input
X = tf.placeholder(name="input", shape=[None, size], dtype=tf.float32)
y = tf.placeholder(name="label", shape=[None, 1], dtype=tf.float32)
# 网络结构
beta = tf.get_variable(name="beta", shape=[size, 1], initializer=tf.glorot_normal_initializer())
bias = tf.get_variable(name="bias", shape=[1], initializer=tf.glorot_normal_initializer())
pred = tf.add(tf.matmul(X, beta), bias, name="output")
# 构建损失
loss = tf.losses.mean_squared_error(y, pred)
# 构建train_op
train_op = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8).minimize(loss)

• 保存为checkpoint模式

checkpoint模式将网络和变量数据分开保存，保存好的模型长这个样子：

|--checkpoint_dir
|    |--checkpoint
|    |--test-model-550.meta
|    |--test-model-550.data-00000-of-00001
|    |--test-model-550.index

checkpoint_dir就是保存时候指定的路径，路径下会生成4个文件。其中.meta文件（其实就是pb格式文件）用来保存模型结构，.data和.index文件用来保存模型中的各种变量，而checkpoint文件里面记录了最新的checkpoint文件以及其它checkpoint文件列表，在inference时可以通过修改这个文件，指定使用哪个model。那么要如何保存呢？

# 只有sess中有变量的值，所以保存模型的操作只能在sess内
checkpoint_dir = "./model_ckpt/"
saver = tf.train.Saver(max_to_keep=1)    # saver 不需要在sess内
with tf.Session() as sess:
    saver.save(sess, checkpoint_dir + "test-model",global_step=i, write_meta_graph=True)

• 保存为pb模式

pb模式保存的模型，只有一个文件，这也是它相比于其他几种方式的优势，简单明了。

pb_dir = "./model_pb/"
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    graph_def = tf.get_default_graph().as_graph_def()
    # 这里是指定要冻结并保存到pb模型中的变量
    var_list = ["input", "label", "beta", "bias", "output"]   # 如果有name_scope，要写全名，如:"name_scope/beta" 
    constant_graph = tf.graph_util.convert_variables_to_constants(sess, graph_def, var_list)
    with tf.gfile.FastGFile(pb_dir + "test-model.pb", mode='wb') as f:
        f.write(constant_graph.SerializeToString())

其实pb模式本质上就是把变量先冻结成常数，然后保存到图结构中。这样就可以直接加载图结构和“参数”了。

• 保存为saved_model模式

虽然saved_model也支持模型加载，并进行迁移学习。可是不得不说saved_model几乎就是为了部署而生的，因为依靠tf.Serving部署模型时要求模型格式必须是saved_model格式。除此以外saved_model还有另外一个优点就是可以跨语言读取，所以本文也介绍一下这种模式的保存于加载。本文样例的保存在参数设置上会考虑到方便部署。保存好的saved_model结构长这个样子：

|--saved_model_dir
|    |--1
|        |--saved_model.pb
|        |--variables
|            |--variables.data-00000-of-00001
|            |--variables.index

保存时需要将保存路径精确到"saved_model_dir/1/ "，会在下面生成一个pb文件，以及一个variables文件夹。其中“1”文件夹是表示版本的文件夹，应该是一个整数。人为设定这个“版本文件夹”的原因是，在模型部署的时候需要将模型位置精确到saved_model_dir，tf.Serving会在saved_model_dir下搜索版本号最大的路径下的模型进行服务。

version = "1/"
saved_model_dir = "./saved_model/test-model-dir/"
builder = tf.saved_model.builder.SavedModelBuilder(saved_model_dir + version)

# 构建 signature
signature = tf.saved_model.signature_def_utils.build_signature_def(
        # 获取输入输出的信息（shape,dtype等），在部署服务后请求带来的数据会喂到inputs中，服务吐的结果会以outputs的形式返回
        inputs={"input": tf.saved_model.utils.build_tensor_info(X)},          # 获取输入tensor的信息，这个字典可以有多个key-value对
        outputs={"output": tf.saved_model.utils.build_tensor_info(pred)},     # 获取输出tensor的信息，这个字典可以有多个key-value对
        method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME    # 就是'tensorflow/serving/predict'
)

# 保存到 saved_model
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    builder.add_meta_graph_and_variables(sess, 
        tags=[tf.saved_model.tag_constants.SERVING],         # 如果用来部署，就这样写。否则可以写其他，如["test-model"]
        signature_def_map={"serving_default": signature},    # 如果用来部署，字典的key必须是"serving_default"。否则可以写其他
    )
    builder.save()

在保存之前需要构建一个signature，用来构造signature的build_signature_def函数有三个参数：inputs、outputs、method_name。其中inputs和outputs分别用来获取输入输出向量的信息，在部署服务后来的数据会喂到inputs中，服务吐的结果会以outputs的形式返回；而method_name如果用来部署模型的话需要设置为"tensorflow/serving/predict", "tensorflow/serving/classify", "tensorflow/serving/regress" 中的一个。如果不是用来服务，就可以写一个其他的。

在保存的时候，除了刚刚构建的signature，还需要提供一个tags 参数，如果用来部署的话需要填[tf.saved_model.tag_constants.SERVING]，否则可以填其他。另外如果用来部署模型的话，signature_def_map的key必须是"serving_default"。

1.3.2 加载

• checkpoint加载（略烦）

checkpoint模式的网络结构和变量是分来保存的，加载的时候也需要分别加载。而网络结构部分你有两种选择：1. 加载.meta文件中的结构， 2. 手动重新写一遍原样结构。

我们先说后一个，如果你不光有模型文件，还有源码，可以把源码构建模型那部分复制过来，然后只加载变量就好，这是手动重新搭建网络结构：

import tensorflow as tf
size = 10
# 构建input
X = ... # 同前面
y = ... # 同前面
# 网络结构
beta = ... # 同前面
bias = ... # 同前面
pred = ... # 同前面

然后加载变量：

# 假设这是一个batch_size=8的batch
feed_X, feed_y = ...
# 用加载出来的参数，跑一下pred
saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint('./model_ckpt/'))    # 加载模型中的变量
    # sess.run(tf.global_variables_initializer())    # 重新初始化一下参数
    print(sess.run(pred, feed_dict={X:feed_X}))

所以手动构建网络结构后，只需要saver.restore一下，就可以加载模型中的参数。另外，通过checkpoint这种模式加载进来的变量，依然是变量，而且是trainable=True的。如果无法手动构建网络呢，就需要从.meta文件里导入网络结构了。

# 不手动构建，从文件中加载网络结构
import numpy as np
import tensorflow as tf
size = 10
# 加载网络
saver=tf.train.import_meta_graph('./model_ckpt/test-model-0.meta')

网络结构已经加载进来了，需要通过name获取tensor

# 假设这是一个batch
feed_X, feed_y = ...
# 下面我们来跑一下 pred
with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint('./model_ckpt/'))  # 加载模型变量
    graph = tf.get_default_graph()
    X = graph.get_tensor_by_name("input:0")        # 根据tensor名字获取tensor变量
    pred = graph.get_tensor_by_name("output:0")    # 根据tensor名字获取tensor变量
    # sess.run(tf.global_variables_initializer())  # 是否重新初始化变量
    print(sess.run(pred, feed_dict={X:feed_X}))

• pb模式加载

相比之下，pb模式的加载旧没那么复杂，因为他的网络结构和数据是存在一起的。

# 直接从pb获取tensor
pb_dir = "./model_pb/"
with tf.gfile.FastGFile(pb_dir + "test-model.pb", "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())    # 从pb文件中导入信息
    # 从网络中通过tensor的name获取为变量
    X, pred = tf.import_graph_def(graph_def, return_elements=["input:0", "output:0"])

现在我们就已经有了X和pred，下面来跑一个pred吧

# 假设这是一个batch
feed_X, feed_y = ...
# 跑一下 pred
with tf.Session() as sess:
    # sess.run(tf.global_variables_initializer())
    print(sess.run(pred, feed_dict={X:feed_X}))

就这么简单！从pb中获取进来的“变量”就可以直接用。为什么我要给变量两个字打上引号呢？因为在pb模型里保存的其实是常量了，此时的“beta:0”和"bias:0"已经不再是variable，而是constant。这带来一个好处：读取模型中的tensor可以在Session外进行。相比之下checkpoint只能在Session内读取模型，对Fine-tune来说就比较麻烦。

• saved_model模式加载

先看一下直接通过tensor的name获取变量的加载方式：

# 假设这是一个batch
feed_X, feed_y = ...

saved_model_dir = "./saved_model/1/"
with tf.Session() as sess:
    # tf.saved_model.tag_constants.SERVING == "serve"，这里load时的tags需要和保存时的tags一致
    meta_graph_def = tf.saved_model.loader.load(sess, tags=["serve"], export_dir=saved_model_dir)
    graph = tf.get_default_graph()
    X = graph.get_tensor_by_name("input:0")
    pred = graph.get_tensor_by_name("output:0")
    # sess.run(tf.global_variables_initializer())
    print(sess.run(pred, feed_dict={X:feed_X}))

这里和checkpoint的加载过程很相似，先一个load过程，然后get_tensor_by_name。这需要我们事先知道tensor的name。如果有了signature的信息就不一样了，saved_model可以通过前面提到的signature_def_map的方法获取tensor。

# 假设这是一个batch
feed_X, feed_y = ...

saved_model_dir = "./saved_model/1/"
with tf.Session() as sess:
    # tf.saved_model.tag_constants.SERVING == "serve"，这里load时的tags需要和保存时的tags一致
    meta_graph_def = tf.saved_model.loader.load(sess, tags=["serve"], export_dir=saved_model_dir)
    signature = meta_graph_def.signature_def
    # print(signature)    # signature 内包含了保存模型时，signature_def_map 的信息
    X = signature["serving_default"].inputs["input"].name
    pred = signature["serving_default"].outputs["output"].name
    print(sess.run(pred, feed_dict={X:feed_X}))

这时即使我们没有源码，也可以通过print(signature)获知关于tensor的信息。

1.3.3 Fine-tune

我们尝试用前面的模型结果作为特征通过一元罗辑回归去预测z，这样新的网络结构就是这样：

import numpy as np
import tensorflow as tf

# 加载模型部分，直接从pb获取X和pred
pb_dir = "./model_pb/"
with tf.gfile.FastGFile(pb_dir + "test-model.pb", "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    X, pred = tf.import_graph_def(graph_def, return_elements=["input:0", "output:0"])

# 下面是 Fine-tune 部分
# 新的 label
z = tf.placeholder(name="new_label", shape=[None, 1], dtype=tf.float32)
# 新的参数
new_beta = tf.get_variable(name="new_beta", shape=[1], initializer=tf.glorot_normal_initializer())
new_bias = tf.get_variable(name="new_bias", shape=[1], initializer=tf.glorot_normal_initializer())
# 一元罗辑回归，通过pred去预测z
new_pred = tf.sigmoid(new_beta * pred + new_beta)    # 这种变量不写name的习惯是不好的哦

# 下面是构建模型的损失函数以及train_op
# log_loss
new_loss = tf.reduce_mean(tf.losses.log_loss(predictions=new_pred, labels=z))
# train_op
train_op = tf.train.AdamOptimizer(learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8).minimize(new_loss)

但是这里存在一个问题，就是只能通过name获取节点。比如这里的new_pred就没有name，那我想要基于这个新模型再次进行Fine-tune的时候，就不能获取这个new_pred，就无法进行Fine-tune。下

# 假设这是一个batch
feed_X, feed_z = ...

# 跑一下 new_pred 之后train一个step，在看看 new_pred 有没有改变
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(new_pred, feed_dict={X:feed_X}))
    sess.run(train_op,  feed_dict={X:feed_X, z:feed_z})
    print(sess.run(new_pred, feed_dict={X:feed_X}))

这里补充一下：通过pb模式导入进来的参数其实是constants，所以在Fine-tune的时候不会变化，而通过checkpoint模式导入进来的参数是variables，在后续Fine-tune的时候是会发生变化的。具体让不让他trainable就看你的实际需要了。

1.4 通过tf.Serving+Docker部署

1.4.1 导出saved_model

只有saved_model形式的模型可以用来部署。假设我们的网络结构采用前文“模型保存、加载与Fine-tune”那一节中的线性回归的结构，保存好的saved_model结构需要长这个样子：

|--saved_model_dir
|    |--version
|        |--saved_model.pb
|        |--variables
|            |--variables.data-00000-of-00001
|            |--variables.index

其中version是版本（数字），在模型保存的时候需要将路径精确到version，在进行部署的时候，指定模型路径只需要精确到saved_model_dir，tf.serving会自动选择版本最大的模型进行服务。

1.4.2 Docker环境搭建

Docker是做什么的呢？说白了就是把一个特定的环境封装起来，如果我们要部署不同的模型，他们有的是需要python2有的需要python3，有的用tf1.x有的用tf2.x，那就需要好几套环境来进行部署，全都装在一个机器里太乱了。所以干脆搞一个集装箱式的东西，每一个集装箱里面装了一套环境（比如py2+tf1.x，或者py3+tf2.x），集装箱里面的环境与外界是独立的，然后把模型部署到这个集装箱里，整个机器的环境就不乱了。

Step1:安装Docker

如何安装Docker不是本篇重点，网上已有的教程比我写的好，我就负责贴转送门。MAC可以直接在这里点击下载Docker并安装。Linux可以这样安装。最后安装完成后，查看一下是否安装成功

$ docker --version
Docker version 19.03.8, build afacb8b

Step2:拉取 tf.Serving镜像

前面提到Docker就是把不同的环境封装成一个集装箱，而一个集装箱里有什么，我们需要一张类似图纸的东西。有了这个图纸，就可以按照这个图纸成产出一个又一个的集装箱。这个图纸学名叫“镜像”，而一个具体的集装箱学名叫“容器”（就好比编程中类和实例的概念）。安装好Docker之后我们来看一下现在有哪些镜像

# 查看现有镜像命令
docker images
# 输出
REPOSITORY          TAG                 IMAGE ID            CREATED             SIZE

我们还没有任何镜像，不过Google官方已经做好了很多镜像，我们可以直接下载合适的版本。官方所有的Tensorflow Docker都在这个Docker Hub代码库中。每个镜像都有自己的“标记”和“变体”，下载时需要根据这个标记和变体来确定我们要下载哪一个版本的镜像。标记有四种

• latest：TensorFlow CPU 二进制映像的最新版本。（默认版本）
• nightly：TensorFlow 映像的每夜版。（不稳定）
• version：指定 TensorFlow 二进制映像的版本，例如：2.1.0
• devel：TensorFlow master开发环境的每夜版。包含 TensorFlow 源代码。

在上面四个标记（tag）之外还有三个变体：

• tag-gpu：支持 GPU 的指定标记版本。
• tag-py3：支持 Python 3 的指定标记版本。
• tag-jupyter：带有 Jupyter 的指定标记版本（包含 TensorFlow 教程笔记本）

比如下面的下载命令都是合法的

docker pull tensorflow/tensorflow                     # latest stable release
docker pull tensorflow/tensorflow:devel-gpu           # nightly dev release w/ GPU support
docker pull tensorflow/tensorflow:latest-gpu-jupyter  # latest release w/ GPU support and Jupyter

这里我们选择tensorflow 1.15.0版本的镜像

docker pull tensorflow/serving:1.15.0

在看看我们有什么镜像

# 查看现有镜像命令
docker images
# 输出
REPOSITORY           TAG                 IMAGE ID            CREATED             SIZE
tensorflow/serving   1.15.0              13090fcf0f63        5 months ago        218MB

现在我们已经成功下载了一个tf.serving的Docker镜像。另外补充一下，如果下载错了的话，可以这样删除镜像：

docker rmi tensorflow/serving:1.15.0

docker rmi后面跟的是镜像的名字+TAG。另外关于容器各种操作可以参考这里。

1.4.3 部署

只要将saved_model导出来，配置好了合适的signature、signature_def_map以及tags。就可以通过Docker进行部署了。这里给出两种简单的进行模型部署的命令：

第一种：

docker run -p 8500:8500 
    --mount type=bind,source=/Users/coreyzhong/workspace/tensorflow/saved_model/,target=/models/test-model 
    -t tensorflow/serving:1.15.0 
    -e MODEL_NAME=test-model --model_base_path=/models/test-model/ &

其中-p 8501:8501表示服务端口；-t表示指定哪一个镜像（图纸）；--mount后面跟的这一大长串中间不能有空格，其中type=bind是选择挂在模式，source指的是模型在机器种储存的路径（必须是绝对路径），target指的是模型在容器中储存的路径（放在集装箱的哪里）。

第二种：

docker run -t --rm -p 8501:8501 -v /Users/coreyzhong/workspace/tensorflow/saved_model/:/models/test-model 
  -e MODEL_NAME=test-model tensorflow/serving:1.15.0 &

其中-p 8501:8501表示服务端口；-v path1:path2中path1指的是模型在机器种储存的路径（必须是绝对路径），path2指的是模型在容器中储存的路径（放在集装箱的哪里）；MODEL_NAME指的是给模型起个名，最后tensorflow/serving:1.15.0是指用哪一个镜像来生成容器。

1.4.4 请求

请求的数据格式应该是json格式，具体形式应该和保存模型时signature的inputs相匹配。

# 假设这是一个请求的样本，包含两个样本：[1,2,3,4,5,6,7,8,9,10]和[1,1,1,1,1,1,1,1,1,1]
data='{"inputs":[[1,2,3,4,5,6,7,8,9,10],[1,1,1,1,1,1,1,1,1,1]]}'

在本地新开一个控制台窗口，并通过curl发送请求

# 通过 curl 发送请求
curl http://localhost:8501/v1/models/test-model:predict -X POST -d "${data}"

# 输出
{
    "outputs": [
        [
            -13.8265877
        ],
        [
            -3.47753382
        ]
    ]
}

这样一个模型就部署好了！

1.5 分布式训练

对于数据量较大的时候，通过分布式训练可以加速训练。相比于单机单卡、单机多卡只需要用with tf.device('/gpu:0')来指定GPU进行计算的情况，分布式训练因为涉及到多台机器之间的分工交互，所以更麻烦一些。本章简单介绍了多机(单卡/多卡不重要)情况下的分布式Tensorflow训练方法。

对于分布式训练与单机训练主要有两个不同：1. 如何开始训练；2. 训练时如何进行分工。分别会在下面两节进行介绍。

1.5.1 确认彼此

单机训练直接可以通过一个脚本就告诉机器“我要开始训练啦”就可以，但是对于分布式训练而言，多台机器需要互相通信，就需要先“见个面认识一下”。就需要给每一台机器一个“名单”，让他去找其他机器。这个“名单”就是所谓的ClusterSpec，让他去找其他机器就是说每一台机器都要运行一次脚本。

下面我们来举一个例子，假设我们用本地机器的两个端口来模拟集群中的两个机器，两个机器的工作内容都是简单的print一句话。首先写两个脚本，第一个脚本长这样

import tensorflow as tf

# 每台机器要做的内容（为了简化，不训练了，只print一下）
c = tf.constant("Hello from server1")

# 集群的名单
cluster = tf.train.ClusterSpec({"local":["localhost:2222", "localhost:2223"]})
# 服务的声明，同时告诉这台机器他是名单中的谁
server = tf.train.Server(cluster, job_name="local", task_index=0)
# 以server模式打开会话环境
sess = tf.Session(server.target, config=tf.ConfigProto(log_device_placement=True))
print(sess.run(c))
server.join()

然后第二个脚本长这样：

import tensorflow as tf

# 每台机器要做的内容（为了简化，不训练了，只print一下）
c = tf.constant("Hello from server2")

# 集群的名单
cluster = tf.train.ClusterSpec({"local":["localhost:2222", "localhost:2223"]})
# 服务的声明，同时告诉这台机器他是名单中的谁
server = tf.train.Server(cluster, job_name="local", task_index=1)
# 以server模式打开会话环境
sess = tf.Session(server.target, config=tf.ConfigProto(log_device_placement=True))
print(sess.run(c))
server.join()

我们来简单说明一下脚本中的内容。这两个脚本其实长的差不多，都是拿着同一个“名单”，即

# 声明集群的“名单”
cluster = tf.train.ClusterSpec({"local":["localhost:2222", "localhost:2223"]})

不同之处只是在创建Server的时候，指定了不同的index，相当于告诉他名单里哪一个名字是自己。其实原理上就是在每一台机器上起一个服务，然后通过这个服务和名单来实现通信。

# 第一个脚本的服务
server = tf.train.Server(cluster, job_name="local", task_index=0)
# 第二个脚本的服务
server = tf.train.Server(cluster, job_name="local", task_index=1)

现在有两个脚本了（对于多机情况，这两个脚本是分别放在不同机器上的，但是本例使用单机的两个端口模仿多机，所以两个脚本可以放在一起）。然后我们让这个“集群”启动起来吧！首先打开一个命令行窗口，在该路径下运行第一个脚本：

# 运行第一台机器（控制台窗口）
$ python3 server1.py

# 输出内容
# 此处省略 N 行内容
2020-04-24 14:58:58.841179: I tensorflow/core/distributed_runtime/master.cc:268] CreateSession still waiting for response from worker: /job:local/replica:0/task:1
2020-04-24 14:59:08.844255: I tensorflow/core/distributed_runtime/master.cc:268] CreateSession still waiting for response from worker: /job:local/replica:0/task:1

忽略WARNING部分，命令行中不断输出内容CreateSession still waiting for response from worker表示这个服务正在等待集群中其他机器，毕竟我们还没有让第二台机器加入进来。下面我们重新打开一个命令行窗口（表示另一台机器），并在目录下启动另一个脚本：

# 运行第二台机器（控制台窗口）
$ python3 server2.py

# 输出内容
# 此处省略 N 行内容
Const: (Const): /job:local/replica:0/task:0/device:CPU:0
2020-04-24 15:02:27.653508: I tensorflow/core/common_runtime/placer.cc:54] Const: (Const): /job:local/replica:0/task:0/device:CPU:0
b'Hello from server2'

我们看到当第二个脚本开始运行时，集群中所有（两台）机器都到齐了，于是就开始工作了。第二台机器直接print出了内容b'Hello from server2'。同时此时第一台机器也开始了工作

# 第二个台机器（控制台窗口）加入到集群之后，第一台机器的输出
Const: (Const): /job:local/replica:0/task:0/device:CPU:0
2020-04-24 15:02:28.732132: I tensorflow/core/common_runtime/placer.cc:54] Const: (Const): /job:local/replica:0/task:0/device:CPU:0
b'Hello from server1'

综上，对于分布式训练来说，第一步就是每一个机器都应该有一个脚本；第二 步给每台机器一个相同的“名单”，也就是ClusterSpec；第三步在每台机器上分别运行脚本，起服务；最后多台机器之间就可以通信了。

1.5.2 密切配合

前一节介绍了集群之间的机器如何相互确认，并一起开始工作的。实际上在复杂的训练过程中会更复杂，我们要为每台机器分配不同的工作，一般会分成ps机和worker机。其中ps机负责保存网络参数、汇总梯度值、更新网络参数，而worker机主要负责正向传导和反向计算梯度。这时在创建ClusterSpec的时候就需要这样做

# 通常将机器分工为ps和worker，不过可以根据实际情况灵活分工。
# 只是在编写代码时明确每种分工的机器要做什么事情就可以
tf.train.ClusterSpec({
    "ps":["localhost:2222"],    # 用来保存、更新参数的机器
    "worker":["localhost:2223", "localhost:2224"]    # 用来正向传播、反向计算梯度的机器
})

本例中仍然采用本机的三个端口模拟三台机器。ClusterSpec的参数字典的key为集群分工的名称，value为该分工下的机器列表。

• 异步分布式训练

我们还是据一个简单的DNN来分类MNIST数据集的例子，脚本应该长这样：

# 异步分布式训练
#coding=utf-8
import time
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data    # 数据的获取不是本章重点，这里直接导入

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("job_name", "worker", "ps or worker")
tf.app.flags.DEFINE_integer("task_id", 0, "Task ID of the worker/ps running the train")
tf.app.flags.DEFINE_string("ps_hosts", "localhost:2222", "ps机")
tf.app.flags.DEFINE_string("worker_hosts", "localhost:2223,localhost:2224", "worker机，用逗号隔开")

# 全局变量
MODEL_DIR = "./distribute_model_ckpt/"
DATA_DIR = "./data/mnist/"
BATCH_SIZE = 32


# main函数
def main(self):
    # ==========  STEP1: 读取数据  ========== #
    mnist = input_data.read_data_sets(DATA_DIR, one_hot=True, source_url='http://yann.lecun.com/exdb/mnist/')    # 读取数据

    # ==========  STEP2: 声明集群  ========== #
    # 构建集群ClusterSpec和服务声明
    ps_hosts = FLAGS.ps_hosts.split(",")
    worker_hosts = FLAGS.worker_hosts.split(",")
    cluster = tf.train.ClusterSpec({"ps":ps_hosts, "worker":worker_hosts})    # 构建集群名单
    server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_id)    # 声明服务

    # ==========  STEP3: ps机内容  ========== #
    # 分工，对于ps机器不需要执行训练过程，只需要管理变量。server.join()会一直停在这条语句上。
    if FLAGS.job_name == "ps":
        with tf.device("/cpu:0"):
            server.join()

    # ==========  STEP4: worker机内容  ========== #
    # 下面定义worker机需要进行的操作
    is_chief = (FLAGS.task_id == 0)    # 选取task_id=0的worker机作为chief

    # 通过replica_device_setter函数来指定每一个运算的设备。
    # replica_device_setter会自动将所有参数分配到参数服务器上，将计算分配到当前的worker机上。
    device_setter = tf.train.replica_device_setter(
        worker_device="/job:worker/task:%d" % FLAGS.task_id,
        cluster=cluster)

    # 这一台worker机器需要做的计算内容
    with tf.device(device_setter):
        # 输入数据
        x = tf.placeholder(name="x-input", shape=[None, 28*28], dtype=tf.float32)    # 输入样本像素为28*28
        y_ = tf.placeholder(name="y-input", shape=[None, 10], dtype=tf.float32)      # MNIST是十分类
        # 省略具体网络结构，直接输出y
        y = ... 
        pred = tf.argmax(y, axis=1, name="pred")
        global_step = tf.contrib.framework.get_or_create_global_step()    # 必须手动声明global_step否则会报错
        # 损失和优化
        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=tf.argmax(y_, axis=1))
        loss = tf.reduce_mean(cross_entropy)
        train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss, global_step=global_step)
        if is_chief:
            train_op = tf.no_op()
    
        hooks = [tf.train.StopAtStepHook(last_step=10000)]
        config = tf.ConfigProto(
            allow_soft_placement=True,    # 设置成True，那么当运行设备不满足要求时，会自动分配GPU或者CPU。
            log_device_placement=False,   # 设置为True时，会打印出TensorFlow使用了哪种操作
        )

        # ==========  STEP5: 打开会话  ========== #
        # 对于分布式训练，打开会话时不采用tf.Session()，而采用tf.train.MonitoredTrainingSession()
        # 详情参考：https://www.cnblogs.com/estragon/p/10034511.html
        with tf.train.MonitoredTrainingSession(
                master=server.target,
                is_chief=is_chief,
                checkpoint_dir=MODEL_DIR,
                hooks=hooks,
                save_checkpoint_secs=10,
                config=config) as sess:
            print("session started!")
            start_time = time.time()
            step = 0
        
            while not sess.should_stop():
                xs, ys = mnist.train.next_batch(BATCH_SIZE)    # batch_size=32
                _, loss_value, global_step_value = sess.run([train_op, loss, global_step], feed_dict={x:xs, y_:ys})
                if step > 0 and step % 100 == 0:
                    duration = time.time() - start_time
                    sec_per_batch = duration / global_step_value
                    print("After %d training steps(%d global steps), loss on training batch is %g (%.3f sec/batch)" % (step, global_step_value, loss_value, sec_per_batch))
                step += 1
    

if __name__ == "__main__":
    tf.app.run()

代码虽然比较长，但是整体结构还是很清晰的。结构上分5个步骤：1. 读取数据、2. 声明集群、3. ps机内容、4. worker机内容、5. 打开会话。其中第四步“worker机内容”包含了网络结构的定义，比较复杂。

接下来只需要将脚本放在集群的三个不同机器上，然后分别运行即可，首先运行ps机脚本：

# ps机脚本
$ python3 distribute_train.py --job_name=ps --task_id=0 --ps_hosts=localhost:2222 --worker_hosts=localhost:2223,localhost:2224

然后运行第一个worker机脚本，开始运行之后他会等待worker其他机的加入：

# 第一个worker机
$ python3 distribute_train.py --job_name=worker --task_id=0 --ps_hosts=localhost:2222 --worker_hosts=localhost:2223,localhost:2224

# 这里省略 N 行输出
2020-04-24 17:25:41.174507: I tensorflow/core/distributed_runtime/master.cc:268] CreateSession still waiting for response from worker: /job:worker/replica:0/task:1
2020-04-24 17:25:51.176111: I tensorflow/core/distributed_runtime/master.cc:268] CreateSession still waiting for response from worker: /job:worker/replica:0/task:1

然后运行第二个worker机的脚本：

# 第二个worker机
$ python3 distribute_train.py --job_name=worker --task_id=0 --ps_hosts=localhost:2222 --worker_hosts=localhost:2223,localhost:2224

# 输出
session started!
After 100 training steps(100 global steps), loss on training batch is 1.59204 (0.004 sec/batch)
After 200 training steps(200 global steps), loss on training batch is 1.10218 (0.003 sec/batch)
# 这里省略 N 行输出

• 同步分布式训练

同样是采用DNN进行MNIST数据集的分类任务：

# 同步分布式训练
# 与异步训练相同，省略

# main函数
def main(self):
    # 前面的内容与异步训练相同，省略

    # 这一台worker机器需要做的计算内容
    with tf.device(device_setter):
        # 网络结构部分与异步训练相同，省略

        # **通过tf.train.SyncReplicasOptimizer函数实现函数同步更新**
        opt = tf.train.SyncReplicasOptimizer(
            tf.train.GradientDescentOptimizer(0.01),
            replicas_to_aggregate=n_workers,
            total_num_replicas=n_workers
        )
        sync_replicas_hook = opt.make_session_run_hook(is_chief)
        train_op = opt.minimize(loss, global_step=global_step)
        if is_chief:
            train_op = tf.no_op()
    
        hooks = [sync_replicas_hook, tf.train.StopAtStepHook(last_step=10000)]    # 把同步更新的hook加进来
        config = tf.ConfigProto(
            allow_soft_placement=True,    # 设置成True，那么当运行设备不满足要求时，会自动分配GPU或者CPU。
            log_device_placement=False,   # 设置为True时，会打印出TensorFlow使用了哪种操作
        )

        # 后面的部分与异步训练一样，省略
    

if __name__ == "__main__":
    tf.app.run()

同步分布式训练与异步分布式训练几乎一样，只有两点差别：

• 优化器要用tf.train.SyncReplicasOptimizer代替tf.train.GradientDescentOptimizer
• hooks要将sync_replicas_hook = opt.make_session_run_hook(is_chief)也加进来

其他的都和异步分布式训练一样，这里就不做赘述了。

二、高级封装——Estimator

实现一个tf.Estimator主要分三个部分：input_fn、model_fn、main三个函数。其中input_fn负责处理输入数据、model_fn负责构建网络结构、main来决定要进行什么样的任务（train、eval、earlystop等等）。

2.1 input_fn

在前面TFrecord制作与读取中介绍了read_and_decode函数，其实就和这里的input_fn逻辑是类似的，都是通过tf.data每次调用会产生一个batch的数据。

def input_fn(filenames, batch_size=32, num_epochs=None, perform_shuffle=False):
    def _parse_fn(record):
        features = {
            "label": tf.FixedLenFeature([], tf.int64),
            "image": tf.FixedLenFeature([], tf.string),
        }
        parsed = tf.parse_single_example(record, features)
        # image
        image = tf.decode_raw(parsed["image"], tf.uint8)
        image = tf.reshape(image, [28, 28])
        # label
        label = tf.cast(parsed["label"], tf.int64)
        return {"image": image}, label

    # Extract lines from input files using the Dataset API, can pass one filename or filename list
    dataset = tf.data.TFRecordDataset(filenames).map(_parse_fn, num_parallel_calls=10).prefetch(500000)    # multi-thread pre-process then prefetch

    # Randomizes input using a window of 256 elements (read into memory)
    if perform_shuffle:
        dataset = dataset.shuffle(buffer_size=256)

    # epochs from blending together.
    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size) # Batch size to use

    iterator = dataset.make_one_shot_iterator()
    batch_features, batch_labels = iterator.get_next()
    return batch_features, batch_labels

2.2 model_fn

model_fn是Estimator中最核心，也是最复杂的一个部分，在这里面需要定义网络结构、损失、train_op、评估结果等各种与网路结构有关的内容。下面通过前面“TFrecord数据制作与读取”中的例子：通过简单的DNN网络来预测label来说明（这一段代码虽然长，但是也是结构化的，不要嫌麻烦一个part一个part的看，其实不复杂的）。

def model_fn(features, labels, mode, params):
    # ==========  解析参数部分  ========== #
    learning_rate = params["learning_rate"]

    # ==========  网络结构部分  ========== #
    # input
    X = tf.cast(features["image"], tf.float32, name="input_image")
    X = tf.reshape(X, [-1, 28*28]) / 255
    # 省略网络结构，直接输出y
    y = ...
    pred = tf.nn.softmax(y, name="soft_max")

    
    # ==========  如果是 predict 任务  ========== #
    predictions={"prob": pred}
    export_outputs = {tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: tf.estimator.export.PredictOutput(predictions)}
    # Provide an estimator spec for `ModeKeys.PREDICT`
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions,
                export_outputs=export_outputs)
    

    # ==========  如果是 eval 任务  ========== #
    one_hot_label = tf.one_hot(tf.cast(labels, tf.int32, name="input_label"), depth=10, name="label")
    # 构建损失
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y, labels=one_hot_label))
    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(tf.math.argmax(one_hot_label, axis=1), tf.math.argmax(pred, axis=1))
    }
    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions,
                loss=loss,
                eval_metric_ops=eval_metric_ops)
    

    # ==========  如果是 train 任务  ========== #
    # 构建train_op
    train_op = tf.train.AdamOptimizer(learning_rate=learning_rate, beta1=0.9, beta2=0.999, epsilon=1e-8).minimize(loss, global_step=tf.train.get_global_step())
    # Provide an estimator spec for `ModeKeys.TRAIN` modes
    if mode == tf.estimator.ModeKeys.TRAIN:
        return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions,
                loss=loss,
                train_op=train_op)

介绍一下model_fn的结构：

• Part1：解析参数部分，本例中以learning_rate为例，展示如何通过param来将参数传递进来，其他参数为了简便，直接用了数值型。
• Part2：网络结构部分。这部分只是负责构建网络结构，从input到pred，不涉及label部分，所以不要把对labels的处理写在这里，因为如果在predict任务中，可能没有label的数据，就会报错。（在这里其实是支持通过tf.keras来构造网络结构，关于tf.keras的用法我在下一章会介绍。
• Part3：predict任务部分。如果任务目的是predict，那么可以直接通过网络结构计算pred，不需要其他操作。设置好export_outputs，并以tf.estimator.EstimatorSpec形式返回即可。
• Part4：eval任务部分。如果是eval任务，除了网络结构以外还需要计算此时的损失、正确率等指标，所以对于loss的定义要放在这一部分。同时设置好评价指标eval_metric_ops，并以tf.estimator.EstimatorSpec形式返回。
• Part5：train任务部分。最后如果是train任务，除了网络结构、loss，还需要优化器、学习率等内容，所以定义train_op的部分在这里进行。最后以tf.estimator.EstimatorSpec形式返回。

2.3 main

main这部分管的就是，我要怎么用这个模型。

def main():
    # ==========  准备参数 ========== #
    task_type = "train"
    model_params = {
        "learning_rate": 0.001,
    }

    # ==========  构建Estimator  ========== #
    config = tf.estimator.RunConfig().replace(
        session_config=tf.ConfigProto(device_count={'GPU': 0, 'CPU': 1}),
        log_step_count_steps=100,
        save_summary_steps=100,
        save_checkpoints_secs=None,
        save_checkpoints_steps=500,
        keep_checkpoint_max=1
    )
    estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir="./model_ckpt/", params=model_params, config=config)

    # ==========  执行任务  ========== #
    if task_type == "train":
        # early_stop_hook 是控制模型早停的控件，下面两个分别是 tf 1.x 和 tf 2.x 的写法
        # early_stop_hook = tf.contrib.estimator.stop_if_no_increase_hook(estimator, metric_name="accuracy",
        early_stop_hook=tf.estimator.experimental.stop_if_no_increase_hook(estimator, metric_name="accuracy", max_steps_without_increase=1000, min_steps=500)
        train_spec = tf.estimator.TrainSpec(input_fn=lambda: input_fn(tr_files, num_epochs=10, batch_size=32), hooks=[early_stop_hook])
        eval_spec = tf.estimator.EvalSpec(input_fn=lambda: input_fn(va_files, num_epochs=1, batch_size=32), steps=None, start_delay_secs=1000, throttle_secs=1)
        tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
    elif task_type == "eval":
        estimator.evaluate(input_fn=lambda: input_fn(va_files, num_epochs=1, batch_size=32))
    elif task_type == "infer":
        preds = estimator.predict(input_fn=lambda: input_fn(te_files, num_epochs=1, batch_size=32), predict_keys="prob")
        with open("./pred.txt", "w") as fo:
            for prob in preds:
                fo.write("%fn" % (np.argmax(prob['prob'])))
    if task_type == "export":
        feature_spec = {
            "image": tf.placeholder(dtype=tf.float32, shape=[None, 28, 28], name="image"),
        }
        serving_input_receiver_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(feature_spec)
        estimator.export_savedmodel("./saved_model/", serving_input_receiver_fn)

其实main中主要做三件事：1. 通过tf.estimator.RunConfig()配置构建Estimator对象；2. 初始化estimator（model_dir如果非空则自动热启动）；3. 执行train/eval/infer/export任务。

• train任务中初始化好TrainSpec和EvalSpec之后可以直接调用tf.estimator.train。也可以使用train_and_evaluate来一边训练一边输出验证集效果。hook可以看作是在训练验证基础上可以实现其他复杂功能的“插件”，比如本例中的early_stop，其他功能还包括热启动、Fine-tune等等，关于hook的用法比较复杂。
• eval任务输出的就是在model_fn函数中eval_metric_ops定义的指标。
• infer任务就是调用estimator.predict获取在model_fn中定义的export_outputs作为预测值。
• export就是将定义Estimator时候模型路径model_dir下的模型导出为saved_model。另外feature_spec指的是一个请求过来所带的数据应该长什么样，对应了model_fn里面的features（即features["image"]），所以这里feature_spec用的是字典的形式，建议model_fn中的features也用字典形式，哪怕是只有一个元素。

最后，直接跑main函数，或者通过tf.app.run()来运行脚本都可以：

# 直接运行 main 函数
main()

# 通过 tf.app.run() 来运行
if __name__ == "__main__":
    tf.app.run()

2.4 分布式训练

对于单机单卡和单机多卡的情况，可以通过tf.device('/gpu:0')来手动控制，这里介绍一下在多机分布式情况下Estimator如何进行分布式训练。Estimator的分布式训练和低阶API的分布式训练类似，都需要提供一份“集群名单”，并且告诉每一台机器他是名单中的谁，并在每台机器上运行脚本。下面看一个例子

import os
import json
import numpy as np
import tensorflow as tf

FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("job_name", "worker", "chief/ps/worker")
tf.app.flags.DEFINE_integer("task_id", 0, "Task ID of the worker running the train")

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'chief': ["localhost:2221"],
        'ps':  ["localhost:2222"],
        'worker': ["localhost:2223", "localhost:2224"]
    },
    'task': {'type': FLAGS.job_name, 'index': FLAGS.task_id}
})

本例采用本地机的两个端口模拟集群中的两个机器，"cluster"表示集群的“名单”信息。"task"表示该机器的信息，"type"表示该机器的角色，"index"表示该机器是列表中的第几个。tf.Estimator中需要指定一个chief机器，ps机也只是在特定的策略下才需要指定（这一点下文介绍）。

除此之外，只需要在tf.ConfigProto中配置train_distribute就可以了：

strategy = tf.distribute.experimental.ParameterServerStrategy()
config = tf.estimator.RunConfig().replace(
    session_config=tf.ConfigProto(
    ... # 这里配置与前面一样
    train_distribute=strategy
)
estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir="./model_ckpt/", params=model_params, config=config)

接下来只需要在每台机器上运行脚本，就可以完成Esitmator的分布式训练了。实际上可以声明不同的strategy，来实现不同的并行策略：

• tf.distribute.MirroredStrategy：单机多卡情况，每一个GPU都保存变量副本。
• tf.distribute.experimental.CentralStorageStrategy：单机多卡情况，GPU不保存变量副本，变量都保存在CPU上。
• tf.distribute.experimental.MultiWorkerMirroredStrategy ：在所有机器的每台设备上创建模型层中所有变量的副本。它使用CollectiveOps，一个用于集体通信的 TensorFlow 操作，来聚合梯度并使变量保持同步。
• tf.distribute.experimental.TPUStrategy：在TPU上训练模型
• tf.distribute.experimental.ParameterServerStrategy：本例中采用的策略，有专门的ps机负责处理变量和梯度，worker机专门负责训练，计算梯度。所以只有在这种策略下，才需要在os.environ['TF_CONFIG']中设置ps机。
• tf.distribute.OneDeviceStrategy：用单独的设备来训练。

三、高级封装——keras

本章介绍另一种高级封装Keras。Keras的特点就是两个字——简单，不用花时间和脑子去研究各种细节问题。

3.1 贯序结构

最简单的情况就是贯序模型，就是将网络层一层一层堆叠起来，比如DNN、LeNet等。下面通过一个LeNet的例子来展示Keras如何实现贯序模型，我们依然采用MNIST数据集举例：

LeNet-5模型结构

首先假设我们已经读到了数据，对于MNIST数据可以通过官方API直接获取，如果是其他数据可以自行进行数据预处理，由于数据读取内容不是本篇介绍重点，所以不做介绍。

(train_images, train_labels), (valid_images, valid_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images.reshape(-1, 28, 28, 1)
valid_images = valid_images.reshape(-1, 28, 28, 1)
train_images, valid_images = train_images / 255.0, valid_images / 255.0

最后数据的格式为 (n, height, width, channel) ，数据和标签的dtype分别为float和int，Keras相比与低阶API和tf.Estimator相比对于数据type的要求比较友好。

接下来开始构建模型

import tensorflow as tf
from tensorflow.keras.optimizers import SGD

# 构建模型结构
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(6, (5,5), activation='tanh', input_shape=(28,28,1)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Conv2D(16, (5,5), activation='tanh'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(120, activation='tanh'),
    tf.keras.layers.Dense(84, activation='tanh'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 模型编译（告诉模型怎么优化）
model.compile(loss='sparse_categorical_crossentropy',    # 损失函数
             optimizer=SGD(lr=0.05, decay=1e-6, momentum=0.9, nesterov=True),     # 优化器
             metrics=['acc'])    # 评估指标

对于贯序模型，只需要调用tf.keras.models.Sequential()，他的参数是一个由tf.keras.layers组成的列表，就可以确定一个模型的结构，然后再简单通过model.compile()就可以确定模型关于“如何优化”方面的信息。很像sklearn的那样简单易用，没有低阶API那种结构和对话的分离，没有必要维护tensor的name。下面看一些怎么开始训练：

history = model.fit(train_images, train_labels, batch_size=32, epochs=1, verbose=1, shuffle=True, validation_data=(valid_images, valid_labels))

就一句fit就解决了！很sklearn。对于evaluate任务也超简单

model.evaluate(test_images, test_labels, verbose=2)

最后对于predict任务，也和sklearn一样

model.predict(test_images)

3.2 复杂结构

贯序模型对于结构复杂的模型，比如层之间出现了分叉、拼接等操作就无法表示了（比如Inception家族）。但是Keras并没有因此放弃，依然是可以很容易的构建复杂结构的网络的。下面来实现一个下图所示的多塔Inception块：

Inception块结构

假设我们在Previous layer处的输入数据的shape为(256, 256, 3)，该结构用Keras这样实现：

import tensorflow as tf

# input数据接口
input_img = tf.keras.layers.Input(shape=(256, 256, 3))

# 分支0
tower0 = tf.keras.layers.Conv2D(64, (1,1), padding="same", activation="relu")(input_img)
# 分支1
tower1 = tf.keras.layers.Conv2D(64, (1,1), padding="same", activation="relu")(input_img)
tower1 = tf.keras.layers.Conv2D(64, (3,3), padding="same", activation="relu")(tower1)
# 分支2
tower2 = tf.keras.layers.Conv2D(64, (1,1), padding="same", activation="relu")(input_img)
tower2 = tf.keras.layers.Conv2D(64, (5,5), padding="same", activation="relu")(tower2)
# 分支3
tower3 = tf.keras.layers.MaxPooling2D((3,3), strides=(1,1), padding="same")(input_img)
tower3 = tf.keras.layers.Conv2D(64, (1,1), padding="same", activation="relu")(tower3)

# 拼接output
output = tf.keras.layers.concatenate([tower0, tower1, tower2, tower3], axis=1)

# 把前面的计算逻辑，分别指定input和output，并构建成网络
model = tf.keras.models.Model(inputs=input_img, outputs=output)

最后构建网络的步骤中也只需要指定inputs和outouts两个参数，在计算时也依然是从后到前追溯计算的。在tf.keras.layers中定义了很多封装好的层，在复杂的网络，只要用到的都是其中的层，就能通过tf.keras实现。

3.3 保存与加载

前面介绍了低阶API的模型保存与加载，相比之下tf.keras的模型保存加载是非常简单的。先来看看模型的保存和加载：

# 保存
model.save("./model_h5/test-model.h5")
# 加载
new_model = tf.keras.models.load_model("./model_h5/test-model.h5")

这样加载进来的模型就可以直接predict和evaluate，当然也可以看作是热启动继续进行增量训练。

new_model.predict(...)
new_model.evaluate(...)
# 热启动 + 增量训练
new_model.fit(...)

3.4 迁移学习

除了进行增量学习以外，tf.keras还可以实现迁移学习。比如我们将刚刚保存好的模型作为预训练模型，选取它的卷积层+池化层并冻结，然后在后面拼接上新的全链接层。

# 加载 base 模型
base_model = tf.keras.models.load_model("./model_h5/test-model.h5")
# 构建新的网络结构
flatten = base_model.get_layer("flatten").output    # 这个layer的name可以通过 base_model.summary()获取
dense = tf.keras.layers.Dense(256, activation='relu')(flatten)
dense = tf.keras.layers.Dense(128, activation='relu')(dense)
pred = tf.keras.layers.Dense(10, activation='softmax')(dense)
# 将 base_model的输入 -> pred 这段网络结构看作是新的模型
fine_tune_model = tf.keras.models.Model(inputs=base_model.input, outputs=pred)

然后我们来将前面的卷积层冻结，具体冻结哪几层需要我们来手动指定。

# 冻结前面的卷积层
for i, layer in enumerate(fine_tune_model.layers): # 打印各卷积层的名字
    print(i, layer.name)
for layer in fine_tune_model.layers[:6]:
    layer.trainable = False

现在一切准备就绪，最后就是编译网络结构，并且进行Fine-tune：

# 新模型编译
fine_tune_model.compile(loss='sparse_categorical_crossentropy', optimizer=SGD(lr=0.05, decay=1e-6, momentum=0.9, nesterov=True), metrics=['acc'])

# Fine-tune
fine_tune_model.fit(training_images, training_labels, batch_size=32, epochs=1, verbose=1, shuffle=True, validation_data=(test_images, test_labels))

3.5 部署

3.5.1 利用tf.Serving+Docker

前面在低阶API那一章介绍了如何将已经导出为saved_model格式的模型，通过tf.Serving+Docker进行部署。对于tf.keras模型来说当然也可以这样，只需要将模型保存为saved_model形式就可以通过tf.Serving+Docker进行部署了：

# 保存为 h5 模型
# model.save("./model_h5/test-model.h5")
    
# 保存为 saved_model
tf.saved_model.save(model, './saved_model_keras/1')

3.5.2 利用flask搭建服务

当然也可以直接写一个基于flask的后台脚本，接收请求后调用model.predict()，然后返回结果

import numpy as np
import tensorflow as tf
from flask import Flask
from flask import request
import json

app = Flask(__name__)

# 配置tf运行环境 + 加载模型
graph = tf.get_default_graph()
sess = tf.Session()
set_session(sess)
model = tf.keras.models.load_model("./model_h5/test-model.h5")

# 用于服务的接口
@app.route("/predict", methods=["GET","POST"])
def predict():
    data = request.get_json()
    if 'values' not in data:
        return {'result': 'no input'}
    arr = data['values']

    global sess
    global graph
    with graph.as_default():
        set_session(sess)
        res = np.argmax(model.predict(np.array(arr)), axis=1)
    return {'result':res.tolist()}

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

这期间会有一些坑，比如要设置graph，否则会出现graph重复导入的问题；以及要设置Session，否则会出现FailedPreconditionError错误。接下来我们来请求一下试试：

import json
import numpy as np
import tensorflow as tf
from urllib import request

# 通过 urllib.request 进行请求
HOST = "http://127.0.0.1:5000/"
image = training_images[0:5].tolist()    # 这里的training_images的格式就是前面训练时training_images的格式
data = json.dumps({'values': image})     # 构造json格式的数据，这里的"values"作为key必须和服务端的"values"相同
req = request.Request(HOST + "predict", headers={"Content-Type": "application/json"})    # 请求
res = request.urlopen(req, data=bytes(data, encoding="utf-8"))
result = json.loads(res.read())          # 字符串转化为字典
print(result)
# 打印结果： {'result': [5, 0, 4, 1, 9]}

关于flask和web服务的内容不是本文重点，所以就不介绍了。至此，tf.keras的内容就介绍完了。

3.6 并行训练

3.6.1 单机多卡

• 结构并行

所谓结构并行，就是对不同的GPU分别计算网络中不同的结构。只需要通过with tf.device_scope('/gpu:0')的方法来手动控制设备，以实现并行。比如对于前面的Inception块的例子中，完全可以改写成

with tf.device_scope('/gpu:0'):
    tower0 = ...
with tf.device_scope('/gpu:1'):
    tower1 = ...
...

• 数据并行

所谓数据并行，就是不同的GPU都是计算整个网络的每一个节点，只是处理的数据不同。这种并行方式不需要手动控制设备，只需要加上一行代码就可以

model = ...           # 和前面一样
model = tf.keras.utils.multi_gpu_model(model, gpus=2)    # 只需要加上这句
model.compile(...)    # 和前面一样
model.fit(...)        # 和前面一样

另外除了这种方法也可以直接指定哪一块GPU进行计算

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3,5"

3.6.2 集群并行

在集群并行中，和低阶API类似，都需要提供一份“集群名单”以及告诉该机器是集群中的谁。并且在集群中的每一台机器上都运行一次脚本，以启动分布式训练。

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        ‘ps’: ["localhost:2222"],
        'worker': ["localhost:2223", "localhost:2224"]
    },
    'task': {'type': FLAGS.job_name, 'index': FLAGS.task_id}
})

本例采用本地机的两个端口模拟集群中的两个机器，"cluster"表示集群的“名单”信息。"task"表示该机器的信息，"index"表示该机器是"worker"列表中的第几个。

训练部分与单机不同的是，只需先声明一个strategy，并且将构造模型部分放在with strategy.scope():中就可以了。

# 数据准备与前面相同
# 声明分布式 strategy
strategy = tf.distribute.experimental.ParameterServerStrategy()

# 在分布式strategy下构造网络模型
with strategy.scope():
    model = ...           # 与定义一个模型
    model.compile(...)    # 模型编译
# 模型训练
history = model.fit(training_images, training_labels, batch_size=32, epochs=3)
# 保存模型
model.save("./model_h5/dist_model.h5")

以上就是分布式keras训练的方法。实际上可以声明不同的strategy，来实现不同的分布式策略，关于strategy已经在tf.Estimator中介绍过。