TensorFlow 1.11.0如何导出不含优化器变量的SavedModel？

你遇到的问题是：训练模型后导出的SavedModel体积是预期的3倍，磁盘占用和加载内存都远超预期，原因是tf.saved_model.simple_save默认保存了所有全局变量，包括Adam优化器为每个可训练变量创建的两个额外状态变量（一阶矩m、二阶矩v），导致总变量数变为原来的3倍。

解决方法 #

在TensorFlow 1.11.0中，simple_save没有提供直接过滤变量的参数，需要改用tf.saved_model.builder.SavedModelBuilder手动构建SavedModel，明确指定只保存推理必需的模型变量。

修改后的示例代码 #

替换原代码中tf.saved_model.simple_save的部分，改用SavedModelBuilder实现：

import tensorflow as tf
print 'tf.__version__', tf.__version__
! rm -rf tmp
! mkdir -p tmp/test_ckpt/
holder = tf.placeholder(tf.int32, [None, 3], name='fea1')
emb_table = tf.get_variable(
    shape=[int(1e7), 20],
    initializer=tf.initializers.random_uniform(minval=-0.04, maxval=0.04),
    dtype=tf.float32,
    name='emb_table',
    trainable=True)
emb = tf.nn.embedding_lookup(emb_table, holder)
loss = pred = tf.reduce_sum(emb, name='pred')
optimzer = tf.train.AdamOptimizer(0.01).minimize(loss)

sess = tf.Session(config=tf.ConfigProto(
    gpu_options=tf.GPUOptions(allow_growth=True, visible_device_list='0'),
    allow_soft_placement=True))
sess.run(tf.local_variables_initializer())
sess.run(tf.global_variables_initializer())
sess.run(optimzer, feed_dict={holder: [[1,2,3], [4,5,6]]})

# 保存ckpt（仅模型变量）
saver = tf.train.Saver(tf.trainable_variables())
saver.save(sess, 'tmp/test_ckpt/ckpt')

# 使用SavedModelBuilder构建精简版SavedModel
builder = tf.saved_model.builder.SavedModelBuilder('tmp/test_saved')
# 定义推理签名
serving_signature = tf.saved_model.signature_def_utils.predict_signature_def(
    inputs={'fea1': holder},
    outputs={'pred': pred}
)
# 仅保存可训练模型变量，排除优化器状态变量
builder.add_meta_graph_and_variables(
    sess,
    [tf.saved_model.tag_constants.SERVING],
    signature_def_map={
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_signature
    },
    variables_to_save=tf.trainable_variables()  # 核心：指定要保存的变量集合
)
# 完成SavedModel导出
builder.save()

expect_size_in_mb = 1e7 * 20 * 4 / 1024 / 1024
print('expect_size_in_mb', expect_size_in_mb)
print('dump model size:')
! du -sh tmp/*

效果验证 #

运行修改后的代码，导出的tmp/test_saved体积会和tmp/test_ckpt一致（约763MB），加载推理时的内存占用也会符合预期。

原理说明 #

• Adam优化器会为每个可训练变量生成m（一阶矩估计）和v（二阶矩估计）两个变量，用于梯度更新，这些变量对推理没有作用，但会被simple_save默认保存。
• SavedModelBuilder的variables_to_save参数可以指定仅保存tf.trainable_variables()（即模型的可训练参数），从而排除优化器的状态变量，大幅减小模型体积。

内容的提问来源于stack exchange，提问作者zhehao