Tensorflow学习笔记：模型训练数据的保存和恢复的简单实例

#! /usr/bin/env python2
# -*- coding: utf-8 -*-

import tensorflow as tf
import numpy as np
import argparse

'''
保存模型训练后参数的简单实例
'''
print('保存和恢复模型训练后参数的简单实例：')

#创建一个图
my_graph = tf.Graph()

with my_graph.as_default():
	var = tf.Variable(0, name='counter') #一个变量，初始值设置为0，但是要会话执行run才会被赋值
	#创建一个op, 实现var + 2
	step = tf.constant(2)
	newVar = tf.add(var, step)
	update = tf.assign(var, newVar)
	# 启动图后, 变量必须先经过`初始化` (init) op 初始化,
	# 首先必须增加一个`初始化` op 到图中.
	init_op = tf.initialize_all_variables()

	#创建saver来保存模型数据
	saver = tf.train.Saver()

#在会话中运行或测试图
def train_or_test(is_test):
	#创建会话，启动图
	with tf.Session(graph = my_graph) as sess:
		if is_test == False: #如果是训练模型
			print('Train begin...')
			sess.run(init_op) #先运行初始化操作
			print('var = %d' % sess.run(var)) #打印初始值

			#更新var,并打印
			for i in range(5):
				sess.run(update)
				print('[%d] var = %d' % (i, sess.run(var)))
				#保存每次迭代的结果，保存的文件从val_iter-1开始而不是0，这个有点搞不明白，有知道原因的麻烦给个留言 哈哈
				saver.save(sess, './model_data1/val_iter', global_step = i)

			saver.save(sess, './model_data1/val_final')

		else: #如果是测试模型
			print('Test begin...')
			for i in range(5)[1:]:
				iter_data_file = './model_data1/val_iter-' + str(i)
				#恢复每次迭代的结果
				saver.restore(sess, iter_data_file)
				print('[%d] var = %d' % (i, sess.run(var)))

			model_data = tf.train.latest_checkpoint('./model_data1/')
			print(model_data) # ./model_data1/val_final
			saver.restore(sess, model_data)
			print('read final var = %d' % sess.run(var))

#必须定义这个main入口
def main(_):
	train_or_test(ARGS.test)

if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument(
		'-t',
		'--test',
		#type = int,
		default = False,
		action = 'store_true', # 运行 ./model_train1.py -t或--test 则ARGS.test被置为True
		help = 'train: True, test: False.'
		)
  
	#ARGS, unparsed = parser.parse_known_args()
	#tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
	ARGS = parser.parse_args()
	print(ARGS)
	tf.app.run()

'''
命令：(1) ./model_train1.py  训练模型
	  (2) ./model_train1.py -t[--test] 测试模型
'''

'''
保存完 model_data1目录下出现：
	checkpoint      (具有最近检查点列表的协议缓冲区)
	val_final       (包含变量的值)
	val_final.meta  (包含图形结构)
	val_iter-1
	val_iter-1.meta
	...
	val_iter-4
	val_iter-4.meta
'''