#! /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
'''
Tensorflow学习笔记:模型训练数据的保存和恢复的简单实例
最新推荐文章于 2024-04-11 20:25:15 发布