I trained tf model in python, dump saved model in python, load and inference in cpp.
I found that the saved model size in disk is triple than it should be, so as the memory usage when load and inference. That’s because saved model include optimizer related variables. So how to exclude those optimizer related variables when dump saved model in tf?
example notebook code:
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]]})
saver = tf.train.Saver(tf.trainable_variables())
saver.save(sess, 'tmp/test_ckpt/ckpt')
tf.saved_model.simple_save(sess, 'tmp/test_saved',
inputs={'fea1': holder,}, outputs={'pred': pred})
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/*
get output:
tf.version 1.11.0
(‘expect_size_in_mb’, 762.939453125)
dump model size:
763M tmp/test_ckpt
2.3G tmp/test_saved
I want a 763 MB saved model but it is 2.3 GB.
@zhehao.zhang,
Welcome to the Tensorflow Forum!
According to save_model documentation:
include_optimizer If True, save optimizer’s state together.
save_traces (only applies to SavedModel format) When enabled, the SavedModel will store the function traces for each layer. This can be disabled, so that only the configs of each layer are stored. Defaults to True. Disabling this will decrease serialization time and reduce file size, but it requires that all custom layers/models implement a get_config() method.
Can you try to use the model.save() method with the save_traces=False and include_optimizer=False arguments?
Note: We are not supporting Tensorflow1.x version and request you please refer to the migration guide to convert 1.x to 2.x.
Thank you!
Thank you for your help. My work environment is tf 1.11.0 and it seems there is no save_traces or include_optimizer arguments in api of tf.saved_model.builder.SavedModelBuilder. It will be hard for me to upgrade from tf 1 to tf 2.
Is there any way to finish this in tf 1?
In tensorflow 1.11.0, there has no save_traces=False or include_optimizer=False arguments, so I have to exclude optimizer variables by myself.
The method is simple:
First, build your graph, add optimizer, train by samples, and finally you dump a checkpoint, whether save optimizer variables or not does not matter here. Then stop this process.
Second, start a new process, rebuild tensorflow graph (your network), but DO NOT add optimizer into your graph. Then restore the former checkpoint and export as saved model. Since you do not add optimizer variables now, the restore function will ignore these optimizer variables in checkpoint.
example code of new process:
! rm -rf tmp/test_saved
import tensorflow as tf
sess = tf.Session(config=tf.ConfigProto(
gpu_options=tf.GPUOptions(allow_growth=True, visible_device_list='0'),
allow_soft_placement=True))
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_shape',
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)
saver = tf.train.Saver()
saver.restore(sess, 'tmp/test_ckpt_all/ckpt')
tf.saved_model.simple_save(sess, 'tmp/test_saved',
inputs={'fea1': holder,}, outputs={'pred': pred})
! du -sh tmp/*