生成检查点文件(chekpoint file),扩展名.ckpt,tf.train.Saver对象调用Saver.save()生成。包含权重和其他程序定义变量,不包含图结构。另一程序使用,需要重新创建图形结构,告诉TensorFlow如何处理权重。
生成图协议文件(graph proto file),二进制文件,扩展名.pb,tf.tran.write_graph()保存,只包含图形结构,不包含权重,tf.import_graph_def加载图形。
模型存储,建立一个tf.train.Saver()保存变量,指定保存位置,扩展名.ckpt。
神经网络,两个全连接层和一个输出层,训练MNIST数据集,存储训练好的模型。
加载数据、定义模型:
#加载数据
mnist = input_data.read_data_sets("MNIST_data/",one_hot=True)
trX,trY,teX,teY = mnist.train.images,mnist.train.labels,mnist.test.images,mnist.test.labels
X = tf.placeholder("float",[None,784])
Y = tf.placeholder("float",[None,10])
#初始化权重参数
w_h = init_weights([784,625])
w_h2 = init_weights([625,625])
w_o = int_weights([625,10])
#定义权重函数
def init_weights(shape):
return tf.Variable(tf.random_normal(shape,stddev=0.01))
#定义模型
def model(X,w_h,w_h2,w_o,p_keep_input,p_keep_hidden):
#第一个全连接层
X = tf.nn.dropout(X,p_keep_input)
h = tf.nn.relu(tf.matmul(X,w_h))
h = tf.nn.dropout(h,p_keep,hidden)
#第二个全连接层
h2 = tf.nn.relu(tf.matmul(h,w_h2))
h2 = tf.nn.dropout(h2,p_keep_hidden)
return tf.matmul(h2,w_o)#输出预测值
#定义损失函数
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(py_x,Y))
train_op = tf.train.RMSPropOptimizer(0.001,0.9).minimize(cost)
predict_op = tf.argmax(py_x,1)
训练模型及存储模型:
#定义存储路径
ckpt_dir = "./ckpt_dir"
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
定义计数器,为训练轮数计数:
#计数器变量,设置它的trainable=False,不需要被训练
global_step = tf.Variable(0,name='global_step',trainable=False)
定义完所有变量,tf.train.Saver()保存、提取变量,后面定义变量不被存储:
#声明完所有变量,调tf.train.Saver
saver = tf.train.Saver()
#之后变量不被存储
non_storable_variable = tf.Variable(777)
训练模型并存储:
with tf.Session() as sess:
tf.initialize_all_variables().run()
start = global_step.eval() #得到global_stepp初始值
print("Start from:",start)
for i in range(start,100):
#128 batch_size
for start,end in zip(