tensorflow 只恢复部分模型参数的实例_Python

tensorflow 只恢复部分模型参数的实例

2020-05-09 09:58平和之心 Python

今天小编就为大家分享一篇tensorflow 只恢复部分模型参数的实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧

我就废话不多说了，直接上代码吧！

									import tensorflow as tf

									def model_1():

									  with tf.variable_scope("var_a"):

									    a = tf.Variable(initial_value=[1, 2, 3], name="a")

									  vars = [var for var in tf.trainable_variables() if var.name.startswith("var_a")]

									  print(len(vars))

									  return vars

									def model_2():

									  vars1 = model_1()

									  with tf.variable_scope("var_b"):

									    a = tf.Variable(initial_value=[1, 2, 3], name="a")

									  vars2 = [var for var in tf.trainable_variables() if var.name.startswith("var")]

									  print(len(vars2))

									  return vars1

									def pretrain_model1():

									  print("-------- model 1 ------")

									  vars = model_1()

									  with tf.Session() as sess:

									    sess.run(tf.global_variables_initializer())

									    saver = tf.train.Saver()

									    saver.save(sess, "./model.ckpt")

									def train_model2():

									  print("-------- model 2 ------")

									  model1_vars = model_2()

									  with tf.Session() as sess:

									    sess.run(tf.global_variables_initializer())

									    saver = tf.train.Saver(var_list=model1_vars)

									    saver.restore(sess, "./model.ckpt")

									    vars = sess.run([model1_vars])

									    for var in vars:

									      print(var)

									step = 2

									if step == 1:

									  pretrain_model1()

									else:

									  train_model2()