转自 https://zhuanlan.zhihu.com/p/77248274
Inside tf.estimator(2) 使用记录
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
- 安装参考资料:
- 自己的总结(在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())
编辑于 2019-08-22