在tensorflow上加载预训练的vgg-16

时间:2017-06-11 20:44:02

标签: tensorflow vgg-net

我正在尝试使用tensorflow r1.1加载预先训练好的vgg-16网络。该网络提供3个文件:

  • saved_model.pb
  • 变量/ variables.index
  • 变量/ variables.data-00000-的-00001

将变量sess初始化为tf.Session()

我使用以下脚本加载网络并提取一些特定的图层:

vgg_path='./'
model_filename = os.path.join(vgg_path, "saved_model.pb")
export_dir = os.path.join(vgg_path, "variables/")

with gfile.FastGFile(model_filename, 'rb') as f:
    data = compat.as_bytes(f.read())
    sm = saved_model_pb2.SavedModel()
    sm.ParseFromString(data)
    image_input, l7, l4, l3 = tf.import_graph_def(sm.meta_graphs[0].graph_def, 
            name='',return_elements=["image_input:0", "layer7_out:0",
            "layer4_out:0", "layer3_out:0"])

tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, image_input)
tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, l7)
tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, l4)
tf.add_to_collection(tf.GraphKeys.GLOBAL_VARIABLES, l3)

saver = tf.train.Saver(tf.global_variables())
print("load data")
saver.restore(sess, export_dir)

初始化变量saver时,脚本会终止并出现以下错误:

  

TypeError:要保存的变量不是变量:Tensor(" image_input:0",   shape =(?,?,?,3),dtype = float32)

如何修复脚本并恢复预先训练好的vgg网络?

1 个答案:

答案 0 :(得分:1)

由于您有SavedModel,因此可以使用tf.saved_model.loader加载它:

with tf.Session() as sess:
    tf.saved_model.loader.load(sess, ["some_tag"], model_dir)