tensorflow模型save和restore

来源:互联网 发布:电脑网络摄像头 编辑:程序博客网 时间:2024/06/05 08:08

把tensorflow每个阶段训练的模型进行保存,供后面预测或者进一步训练时使用。保存的时候使用saver.save()方法,恢复的时候使用saver.restore方法。详细如下:

初始信息

flags.DEFINE_integer('epochs', 100, '') #FLAGS.epochsflags.DEFINE_integer('display_steps', 10, 'Number of steps to run trainer.') #FLAGS.batch_sizeflags.DEFINE_integer('batch_size', 256, 'Batch size.Must divide evenly into the dataset sizes.')  #FLAGS.batch_sizeflags.DEFINE_string('train_dir', '/data1/Image/image_75', 'Directory to put the training data.')

模型save部分

for epoch in range(FLAGS.epochs):           for batch_i in range(9):                for batch_features, batch_labels in  img_util_getBatch(batch_i, FLAGS.batch_size):                   train_neural_network(sess, optimizer, FLAGS.keep_probability, batch_features, batch_labels)                print 'Epoch {:>2}, CIFAR-10 Batch {}:  '.format(epoch + 1, batch_i)                print_stats(sess, batch_features, batch_labels, cost)#, accuracy)           if epoch%FLAGS.display_steps == 0:             save_path = saver.save(sess, FLAGS.base_dir+'tf/', global_step=epoch, write_meta_graph=True)             print("Model saved in file: %s" % save_path)        saver.export_meta_graph(FLAGS.base_dir+'tf/'+'.meta')        tf.train.write_graph(sess.graph_def, FLAGS.base_dir+'tf/', 'soft.ph',False)

需要注意的地方:
1. global_step的设置意义,在于给文件名字编号,如:

saver.save(sess,'my-model',global_step=10) --> filename:'my-model-10'tf.train.Saver.__init__(varlist=None,reshape=False,shared=False,max_to_keep=5,keep_checkpoint_every_n_hour=100,name=None)

模型restore部分

    image = cv2.imread(path)    x = tf.placeholder(tf.float32, shape=[None, 227, 225, 3], name='x')    y = conv_net(x, 1)    saver = tf.train.Saver()    saver.restore(sess,tf.train.latest_checkpoint('/data1/wzy/deepImage/0919/tf2/'))    feed_dict = {x: [image]}    score = sess.run(y, feed_dict)

图片预测

    image = cv2.imread(image_path)    x = tf.placeholder(tf.float32, shape=[None, 227, 225, 3], name='x')    y = conv_net(x, 1)    sess.run(tf.global_variables_initializer())    saver = tf.train.import_meta_graph('/data1/Image/tf/20.meta')    saver.restore(sess,tf.train.latest_checkpoint('/data1/Image/tf/'))    feed_dict = {x: [image]}    score = sess.run(y, feed_dict)
原创粉丝点击