0. 前言

• 基本使用请参考：Inside tf.estimator（1） 基本使用
• 目标：记录使用 tf.estimator 时碰到的问题。
• 感想：
• 近来在努力搬砖，终于像个真正的算法工程师了。
• 使用的解决方案不一定最优，但都能解决自己的问题。
• 如果有更好的解决方案，请告诉我……
• 本文内容
• tf.keras.layers.BatchNormalization 采坑
• tf.keras.layers.BatchNormalization 不会将修改mean/var/beta等操作自动添加到UPDATE_OPS中。
• 使用 model.updates 也存在问题，需要使用 model.get_updates_for 才行。
• 单机多卡训练/预测/评估
• 通过 MirroredStrategy 实现。
• 导入 fine-tune 模型
• 导入fine-tune模型，不能使用 model.load_weights 实现，真的坑。
• 可以通过 tf.train.init_from_checkpoints / tf.train.Scaffold / tf.estimator.WarmStartSettings 实现。本文介绍后两种。
• 保存验证集误差最小的模型
• 通过 tf.train.SessionRunHook 实现。

1. tf.keras.layers.BatchNormalization 采坑

• 问题描述：
• 使用 tf.keras 构建模型，通过自定义 tf.estimator.EstimatorSpec，构建了 tf.estimator.Estimator 对象。
• 在创建 train_op 时，在optimizer.minimize之前使用了 with tf.control_dependencies(update_ops)。
• 当使用 vgg16 作为backend时，模型能够正常；当使用 resnet50, xception 等作为backend时，效果贼差。
• 发现问题：
• 检查保存下来的ckpt文件时发现，所有bn的mean为0，var为1，也就是说UPDATE_OPS里没有添加BN的更新操作。
• 也就是说 tf.keras.layers.BatchNormalization 没有默认向 UPDATE_OPS 中添加默认操作，而tensorflow其他的bn实现都这么做了……
• 解决：
• 获取所有BN的OPS，添加到tf.GraphKeys.UPDATE_OPS。
• 方法一（参考issue）：
• 又采坑（这个方法不行）：在通过keras的方法获取，即model.updates，是不行的。
• 应该使用 model.get_updates_for().
• 方法二：参考这篇博客，通过OPS的名称寻找所有符合要求的OPS。

ops = tf.get_default_graph().get_operations()
update_ops = [x for x in ops if ("AssignMovingAvg" in x.name and x.type=="AssignSubVariableOp")]
tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_ops)

• 感想
• ………………………………………………………………
• 查了下issue，一年多前就有人问了……大哥，好歹文档里说明下。啊…………………………………………
• 我一直比较抵触使用 tf.keras，但官方就是推荐使用，那就用吧，还能咋办呢。但这个BUG表明，官方的意思是使用全套keras……
• keras的作者还在 issue 19643 中宣称：
• 就是这么设计的。
• 想解决，要么自己手动处理，想要方便就直接用全套 keras API好啦（言下之意，要官方修复是不可能的）……

That is by design. Global collections (and global states in general) are prone to various issues (like most global variables in software engineering) and should be avoided.
...
If you are writing your own custom training loops, you will have to run these updates as part of the call to sess.run(). But do note that tf.keras provides a built-in training loop in the form of model.fit(), as well as primitive methods for building your own custom training loops (in particular model.train_on_batch() and model.predict_on_batch()). If you are using Keras you should generally never have to manually manage your model‘s updates.

2. 单机多卡

• 目标：单机多卡训练/预测/评估。

2.1. 多GPU训练的第一种实现方式

• 最基本、底层的实现方式，不推荐使用。
• 思路：
• 使用 tf.variable_scpe、tf.device等对同一个计算图使用多个gpu创建。
• 分别计算每个GPU中的梯度值，最终汇总求平均，最为最终目标梯度值。
• 实例：官方cifar10示例

2.2. 多GPU训练的第二种方式

• 使用MirroredStrategy实现，强烈推荐。
• 在tf1中，该类位于tensorflow.contrib.distribute.
• tf1.13后可以直接使用tf.distribute.MirroredStrategy来调用，但API有变动……
• 准备工作：安装NCCL
• 必须安装nccl，否则会报错 libnccl.so.2: cannot open shared object file: No such file or directory。
• 相关issue：issue 22899
• 安装参考资料：
• NVIDIA官方教程（推荐）
• 与官方教程类似的中文博客
• 若已经安装了CUDA，但没有安装NCCL，且CUDA版本在9.2以上，可以看这个教程
• 自己的总结（在ubuntu16.04中为cuda9.0安装nccl）：
• 建议直接看文档。
• 首先下载nccl安装文件，地址。（像下载cudnn一样，需要登录，并选择合适的版本）

• 如果CUDA版本高于9.2，则可以下载 O/S agnostic local installer，可以直接下载后解压、将相关文件复制到/usr/local/cuda的对应文件夹中，方便很多，更多请参考上面参考资料中的链接。
• 如果是cuda9，只能下载deb安装包，并下载。
• 下载deb包后，要进行安装：
• 如果是local版，则通过sudo dpkg -i nccl-repo-<version>.deb安装。
• 如果是network版，则通过sudo dpkg -i nvidia-machine-learning-repo-<version>.deb安装。
• 更新APT数据库：sudo apt update
• 通过apt安装libnccl2库，命令在下载nccl文件的网址中有，大概形式就是：sudo apt install libnccl2=2.4.7-1+cuda9.0 libnccl-dev=2.4.7-1+cuda9.0。
• 实现介绍
• 参考资料：Multi-GPU training with Estimators, tf.keras and tf.data
• 在创建tf.estimator.Estimator的tf.estimator.RunConfig中指定train_distribute。
• 其他不需要任何修改……真的方便……
• 参考代码如下：

NUM_GPUS = 2
strategy = tf.contrib.distribute.MirroredStrategy(num_gpus=NUM_GPUS)
config = tf.estimator.RunConfig(train_distribute=strategy)
estimator = tf.keras.estimator.model_to_estimator(model,
                                                  config=config)

2.3. 方法二的一些采坑以及使用记录

• 使用多GPU时调用train/predict/evaluate方法时，init_fn 参数注意事项
• 返回的必须是tf.data.Dataset对象。
• 创建tf.data.Dataset对象的过程必须在init_fn函数中。
• 否则会报类似ValueError: Tensor(xxx) must be from the same graph as Tensor(xxx).的错误
• 参考资料：issue 20784，官方回复是说：One of the key features of estimator is that it performs graph and session management for you.，大概意思是，所有的创建计算图的过程，最好都封装在init_fn和model_fn中，等Estimator调用时在创建计算图。
• 其他：
• 假设创建tf.data.Dataset中有指定batch_size，且确定GPU数量为 num_gpus，则多GPU训练时，数据集实际batch size为batch_size * num_gpus。
• 建议在创建tf.data.Dataset对象时引入prefetch。
• 在未使用MirroredStrategy时，可以在init_fn外创建好tf.data.Dataset，而在init_fn中只创建iterator并调用get_next函数。
• 若使用 slim 构建模型，很可能会报错
• 根据issue 27392，问题可能出在ExponentialMovingAverage操作上。
• 根据issue 20874，当模型中使用 slim.batch_norm 时会与 MirroredStrategy 冲突。
• 个人看法：
• 如果没有 batch_norm，可能模型可以正常运行。但有几个模型是完全不包括batch_norm的……
• 建议使用 tf.keras 构建模型；如果硬要使用 slim 模型，可以将slim.batch_norm改为tf.keras.layers.BatchNormalization。
• slim模型 + tf.estimator + MirroredStrategy 这个组合目前来看应该必须放弃，因为tensorflow的Member已经表态 issue23770 I don‘t think we have any current effort to add distribute strategy support to slim, and it seems unlikely to become a priority due to the general TF move away from contrib.

3. 导入 fine-tune 模型

• 目标：在训练之前，导入部分模型的pre-trained model。

3.1. 使用 slim 搭建模型时

• 基本步骤：
• 第一步：定义init_fn函数，用于构建tf.train.Scaffold对象。
• 函数基本形式：init_fn(scaffold, session)，用于初始化模型参数。
• 构建 init_fn 可以用到 slim.assign_from_checkpoint_fn，细节就不描述了。
• 第二步：将 tf.train.Scaffold 对象，作为 mode == tf.estimator.ModeKeys.PREDICT 时 tf.estimator.EstimatorSpec 对象的初始化参数。
• 采坑，曾经想用 tf.train.SessionRunHook 来实现这个功能：
• 实现方法（最终失败）：构建子类，在 def after_create_session(session, coord) 方法中将finetune模型参数导入，并将该hook添加到tf.estimator.EstimatorSpec中。
• 问题：
• 如果训练中止、再启动时，需要导入训练过程中保存的模型文件，即导入model_dir中最新的模型参数。
• 导入fine-tune模型参数的操作在每次重起训练时都会调用，而且调用顺序在导入model_dir模型参数之后……

3.2 使用 tf.keras 搭建模型时

3.2.1 问题描述：

• 在使用 tf.keras.applications 中的模型都有对应的预训练模型，想通过这些h5文件进进行fine-tune。
• 采坑：不能直接使用 load_weights，原因如下：
• load_weights的最后一步是通过 keras.backend.get_session() 执行 assign 操作，也就是说，会自动创建 tf.Session 对象。
• tf.estimator的目标是自动管理 tf.Session 的生命周期，会创建新的 Session 对象，而不会使用 keras.backend.get_session() 中的对象。
• TensorFlow模型中不能直接使用 h5 文件，所以需要先将h5文件转换为ckpt。
• 在 tf.estimator 中导入fine-tune模型的方法有很多，可以使用以下几种：
• tf.train.init_from_checkpoint：改变的是变量的 initializer，看issue里有个Member提到这事推荐方法，不过我没试过。
• tf.train.Scaffold：定义其中的 init_fn，将scaffold对象传入 tf.estimator.Estimator 的初始化方法或train方法。
• 使用 tf.estimator.WarmStartSettings 对象，导入 tf.estimator.Estimator 初始化方法中。

3.2.2. h5 to ckpt

• 参考keras issue9040，fchollet给出了方法。
• 我自己写的测试样例

import tensorflow as tf
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"]="1"
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
tf.keras.backend.set_session(tf.Session(config=config))
path_to_save_ckpt = ‘/path/to/keras-ckpt‘

model = tf.keras.applications.VGG16()
model_name = ‘vgg16‘

var_list = slim.get_variables_to_restore(include=None, exclude=[‘predictions‘])
saver = tf.train.Saver(var_list)
saver.save(tf.keras.backend.get_session(), 
           os.path.join(path_to_save_ckpt, model_name + ‘.ckpt‘))

3.2.3 tf.train.Scaffold 对象导入 ckpt 文件实现思路

• tf.train.Scaffold对象中init_fn方法的定义为：def init_fn(scaffold, session)。
• 通过 slim.get_variables_to_restore(include=[], exclude=[]) 获取对应的变量。
• 通过 slim.assign_from_checkpoint_fn 获取 def init_fn(session) 方法。
• 该方法可将ckpt文件中tensor和当前模型中实际tensor对应起来。
• var_list 可以是list，也可以是dict。
• dict是 ckpt name to tf.Variable。
• list则表明ckpt文件与当前模型中tensor name是一一对应的，即 {var.op.name: var for var in var_list}
• 有了 init_fn(session) 再自己创建一个 _init_fn(scaffold, session) wrap 一下就行了。

3.2.4 tf.estimator.WarmStartSettings 对象导入 ckpt 文件实现思路：

• 该类就是为了实现部分导入模型参数的功能，该类的注释中有几个简单的使用
• 定义该类时有以下几个参数，现在分别说明：
• ckpt_to_initialize_from：必须填写，ckpt文件路径。
• vars_to_warm_start：需要导入的参数。
• 默认是使用所有 trainable 参数。
• 如果设置为 .*，则表示所有 trainable 参数。
• 如果设置为 [.*]，则表示所有参数（包括non-trainable）。
• 可以设置为[string1, string2]，分别获取。
• 小技巧，如果想要设置以XXX结尾，则使用*xxx[^/]。
• var_name_to_vocab_info：不知道是啥玩意，没仔细研究。
• var_name_to_prev_var_name：如果ckpt name和当前模型tensor name不对应，可以在这里设置。
• 需要注意的是，var_name_to_prev_var_name 不用于过滤参数。
• 换句话说，如果ckpt name和当前模型的tensor name相同时，可以不在 var_name_to_prev_var_name 中进行配置。
• 采坑：如果选择的当前模型参数不存在于ckpt文件中，则会报错。
• 如所有以 Momentum 结尾的参数。
• 解决方法：设置正则表达式，例如设置以kernel或bias结尾的所有变量。
• 举个例子：
• 对于 keras vgg16 网络，若想所有卷基层的tensor name和当前模型的var name相同，但全连接层不相同，则

ws = tf.estimator.WarmStartSettings(
    # ckpt 文件路径
    ckpt_to_initialize_from=os.path.join(ckpt_root_path, ckpt_name),

    # 获取所有block开头、kernel/bias结尾的当前模型var
    # 获取所有original_ranking/rank_fc开头、kernel/bias结尾的当前模型var
    vars_to_warm_start=[‘block.+kernel[^/]‘, 
                        ‘block.+bias[^/]‘,
                        ‘original_ranking/rank_fc.+kernel[^/]‘,
                        ‘original_ranking/rank_fc.+bias[^/]‘],

    # 若当前模型var name与ckpt name不匹配，则在这里进行处理
    # 当前模型 var name to ckpt name
    var_name_to_prev_var_name={
        ‘original_ranking/rank_fc1/bias‘: ‘fc1/bias‘,
        ‘original_ranking/rank_fc1/kernel‘: ‘fc1/kernel‘,
        ‘original_ranking/rank_fc2/bias‘: ‘fc2/bias‘,
        ‘original_ranking/rank_fc2/kernel‘: ‘fc2/kernel‘,
    }
)

4. 保存验证集误差最小的模型

• tf.estimator.Estimator中，基本的训练过程是类似于以下代码

for i in range(args.num_epochs):
    # train
    estimator.train(_train_input_fn,
                    hooks=train_hooks,
                    steps=args.train_steps)

    if i % args.validation_step == 0:
        # val every {validation_step} steps
        estimator.evaluate(_val_input_fn,
                           args.val_steps,
                           hooks=evaluate_hooks)

• 目标：在模型训练过程中，保留验证集误差最小的模型。
• 换句话说：希望根据evaluate的结果保存 loss 最小或 accuracy 最大的模型。
• 实现思路：
• 构建hook，每次evaluate结束后记录loss/accuracy的平均数，如果loss/accuracy比之前记录的对应数值小/大，则保存该模型。
• 难点：由于tf.estimator.Estimator封装了evaluate的过程，无法直接获取loss/accuracy的平均值。
• 解决方案：
• 猜测evaluate中使用了类似 tf.metrics.mean 的操作，用于保存验证集中的平均loss/accuracy。
• 找到 loss/accuracy 平均数对应的tensor name。
• 我使用的方法是，在tensorboard中找到该名称。
• 使用 tf.get_default_graph().get_tensor_by_name(tensor_name) 来获取tf.Tensor对象。
• 构建hook对象，实现上述思路，并将该hook添加到evaluate方法中。
• 源码示例（保存loss最小的模型）：

class EvaluateCheckpointHook(tf.train.SessionRunHook):
    def __init__(self,
                 eval_dir, # 模型保存路径
                 tensor_name=‘mean/value:0‘, # 需要对比的tensor名称
                 ):
        self._cur_min_value = 1e8 # 记录当前最小值
        self._tensor_name = tensor_name
        self._eval_dir = eval_dir

    def begin(self):
        self._saver = tf.train.Saver() # saver 必须在 begin 方法中创建

    def end(self, session):
        target_tensor = tf.get_default_graph().get_tensor_by_name(
            self._tensor_name)  # 获取tensor
        cur_value = session.run(target_tensor)
        if self._cur_min_value > cur_value:
            self._cur_min_value = cur_value
            self._saver.save(session, self._eval_dir,
                             global_step=tf.train.get_global_step())