TensorFlow下网络模型的存储与加载_tensorflow保存和重载网络-优快云博客

本文介绍了一个使用TensorFlow训练MNIST数据集的手写数字识别模型，并详细展示了如何保存及恢复该模型。通过代码实现，读者可以了解到模型训练流程、成本函数定义、优化函数设置以及模型保存的具体操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

     
      #!/usr/bin/env python
     
     
      # 导入mnist数据库
     
     
      from tensorflow.examples.tutorials.mnist 
      import input_data
     
     
      mnist = input_data.read_data_sets(
      "MNIST_data/", one_hot=
      True)
     
     
     
      import tensorflow 
      as tf
     
     
      import os
     
     
     
      # 定义输入变量
     
     
      x = tf.placeholder(tf.float32, [
      None, 
      784])
     
     
     
      # 定义参数
     
     
      W = tf.Variable(tf.zeros([
      784, 
      10]))
     
     
      b = tf.Variable(tf.zeros([
      10]))
     
     
     
      # 定义激励函数
     
     
      y = tf.nn.softmax(tf.matmul(x, W) + b)
     
     
     
      # 定义输出变量
     
     
      y_ = tf.placeholder(tf.float32, [
      None, 
      10])
     
     
     
      # 定义成本函数
     
     
      cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[
      1]))
     
     
     
      # 定义优化函数
     
     
      train_step = tf.train.GradientDescentOptimizer(
      0.5).minimize(cross_entropy)
     
     
     
      # 初始化变量
     
     
      init = tf.global_variables_initializer()
     
     
     
      # 定义会话
     
     
      sess = tf.Session()
     
     
     
      # 运行初始化
     
     
      sess.run(init)
     
     
     
      # 定义模型保存对象
     
     
      saver = tf.train.Saver()
     
     
     
      tf.add_to_collection(
      'x', x)
     
     
      tf.add_to_collection(
      'y', y)
     
     
     
      # 循环训练1000次
     
     
      for i 
      in range(
      1000):
     
     
          batch_xs, batch_ys = mnist.train.next_batch(
      100)
     
     
          sess.run(train_step, feed_dict={x: batch_xs, y_:batch_ys})
     
     
     
      print(
      "训练完成！")
     
     
     
      # 创建模型保存目录
     
     
      model_dir = 
      "mnist_1"
     
     
      model_name = 
      "ckp"
     
     
      if 
      not os.path.exists(model_dir):
     
     
          os.mkdir(model_dir)
     
     
     
      # 保存模型
     
     
      saver.save(sess, os.path.join(model_dir, model_name))
     
     
     
      print(
      "保存模型成功！")

模型使用

模型恢复与使用代码如下：

     
      #!/usr/bin/env python
     
     
      # 导入mnist数据库
     
     
      from tensorflow.examples.tutorials.mnist 
      import input_data
     
     
      mnist=input_data.read_data_sets(
      "MNIST_data",one_hot=
      True)
     
     
     
      import tensorflow 
      as tf
     
     
     
      # 创建会话
     
     
      sess = tf.Session()
     
     
     
      new_saver = tf.train.import_meta_graph(
      'mnist_1/ckp.meta')
     
     
      new_saver.restore(sess, 
      'mnist_1/ckp')
     
     
     
      x = tf.get_collection(
      'x')[
      0]
     
     
      y = tf.get_collection(
      'y')[
      0]
     
     
     
      print(
      "恢复模型成功！")
     
     
     
      # 取出一个测试图片
     
     
      idx = 
      0
     
     
      img = mnist.test.images[idx]
     
     
     
      # 根据模型计算结果
     
     
      ret = sess.run(y, feed_dict = {x : img.reshape(
      1, 
      784)})
     
     
     
      print(
      "计算模型结果成功！")
     
     
     
      # 显示测试结果
     
     
      print(
      "预测结果:%d"%(ret.argmax()))
     
     
      print(
      "实际结果:%d"%(mnist.test.labels[idx].argmax()))
     
     

     
     

     
     

     
     

     
     
      转载请注明出处：
      http://wustmeiming.github.io/2017/01/09/Tensorflow模型保存与使用-导入图方案/