當(dāng)前位置：首頁(yè) > 编程资源 > 编程问答 >内容正文

编程问答

TensorFlow模型保存和加载方法

發(fā)布時(shí)間：2025/3/19 编程问答 33 豆豆

生活随笔收集整理的這篇文章主要介紹了 TensorFlow模型保存和加载方法小編覺(jué)得挺不錯(cuò)的,現(xiàn)在分享給大家,幫大家做個(gè)參考.

TensorFlow模型保存和加載方法

模型保存

import tensorflow as tfw1 = tf.Variable(tf.constant(2.0, shape=[1]), name="w1-name") w2 = tf.Variable(tf.constant(3.0, shape=[1]), name="w2-name") a = tf.placeholder(dtype=tf.float32, name="a-name") b = tf.placeholder(dtype=tf.float32, name="b-name") y = a * w1 + b * w2 init = tf.global_variables_initializer() saver = tf.train.Saver() with tf.Session() as sess: sess.run(init) print(a) # Tensor("a-name:0", dtype=float32) print(b) # Tensor("b-name:0", dtype=float32) print(y) # Tensor("add:0", dtype=float32) print(sess.run(y, feed_dict={a: 10, b: 10})) saver.save(sess, "./model/model.ckpt")

這段代碼中，通過(guò)saver.save函數(shù)將TensorFlow模型保存到了model/model.ckpt文件中，這里代碼中指定路徑為"model/model.ckpt"，也就是保存到了當(dāng)前程序所在文件夾里面的model文件夾中。

TensorFlow模型會(huì)保存在后綴為.ckpt的文件中。保存后在save這個(gè)文件夾中實(shí)際會(huì)出現(xiàn)3個(gè)文件，因?yàn)門(mén)ensorFlow會(huì)將計(jì)算圖的結(jié)構(gòu)和圖上參數(shù)取值分開(kāi)保存。

model.ckpt.meta文件保存了TensorFlow計(jì)算圖的結(jié)構(gòu)，可以理解為神經(jīng)網(wǎng)絡(luò)的網(wǎng)絡(luò)結(jié)構(gòu)
model.ckpt文件保存了TensorFlow程序中每一個(gè)變量的取值
checkpoint文件保存了一個(gè)目錄下所有的模型文件列表

模型加載:只加載變量，但是還是需要重新定義圖結(jié)構(gòu)

import tensorflow as tf# 使用和保存模型代碼中一樣的方式來(lái)聲明變量 # 變量rw1, rw2 不需要進(jìn)行初始化 rw1 = tf.Variable(tf.constant(2.0, shape=[1]), name="w1-name") rw2 = tf.Variable(tf.constant(3.0, shape=[1]), name="w2-name") # 重新定義圖結(jié)構(gòu) result = 10 * rw1 + 10 * rw2 saver = tf.train.Saver() with tf.Session() as sess: saver.restore(sess, "./model/model.ckpt") print(sess.run(result))

tf.train.Saver類也支持在保存和加載時(shí)給變量重命名

import tensorflow as tf# 聲明的變量名稱name與已保存的模型中的變量名稱name不一致 rw1 = tf.Variable(tf.constant(2.0, shape=[1]), name="rw1-name") rw2 = tf.Variable(tf.constant(3.0, shape=[1]), name="rw2-name") # 重新定義圖結(jié)構(gòu) result = 10 * rw1 + 10 * rw2 # 若直接生命Saver類對(duì)象，會(huì)報(bào)錯(cuò)變量找不到 # 使用一個(gè)字典dict重命名變量即可，{"已保存的變量的名稱name": 重命名變量名} # 原來(lái)名稱name為 w1-name 的變量現(xiàn)在加載到變量 rw1（名稱name為 rw1-name）中 saver = tf.train.Saver({"w1-name": rw1, "w2-name": rw2}) with tf.Session() as sess: saver.restore(sess, "./model/model.ckpt") print(sess.run(result))

模型加載: 不需要重新定義圖結(jié)構(gòu)

import tensorflow as tfsaver = tf.train.import_meta_graph("./model/model.ckpt.meta") graph = tf.get_default_graph()# 通過(guò) Tensor 名獲取變量 a = graph.get_tensor_by_name("a-name:0") b = graph.get_tensor_by_name("b-name:0") y = graph.get_tensor_by_name("add:0") with tf.Session() as sess: saver.restore(sess, "./model/model.ckpt") print(sess.run(y, feed_dict={a: 10, b: 10}))

convert_variables_to_constants

# 通過(guò)convert_variables_to_constants函數(shù)將計(jì)算圖中的變量及其取值通過(guò)常量的方式保存于一個(gè)文件中import tensorflow as tf from tensorflow.python.framework import graph_util v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1") v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2") result = v1 + v2 with tf.Session() as sess: sess.run(tf.global_variables_initializer()) # 導(dǎo)出當(dāng)前計(jì)算圖的GraphDef部分，即從輸入層到輸出層的計(jì)算過(guò)程部分 graph_def = tf.get_default_graph().as_graph_def() output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, ['add']) with tf.gfile.GFile("Model/combined_model.pb", 'wb') as f: f.write(output_graph_def.SerializeToString()) # 載入包含變量及其取值的模型 import tensorflow as tf from tensorflow.python.platform import gfile with tf.Session() as sess: model_filename = "Model/combined_model.pb" with gfile.FastGFile(model_filename, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) result = tf.import_graph_def(graph_def, return_elements=["add:0"]) print(sess.run(result)) # [array([ 3.], dtype=float32)]

TensorFlow 模型保存/載入的兩種方法

我們?cè)谏暇€使用一個(gè)算法模型的時(shí)候，首先必須將已經(jīng)訓(xùn)練好的模型保存下來(lái)。tensorflow保存模型的方式與sklearn不太一樣，sklearn很直接，一個(gè)sklearn.externals.joblib的dump與load方法就可以保存與載入使用。而tensorflow由于有g(shù)raph,operation 這些概念，保存與載入模型稍顯麻煩。

一、基本方法

網(wǎng)上搜索tensorflow模型保存，搜到的大多是基本的方法。即

保存

定義變量

使用saver.save()方法保存

載入

定義變量

使用saver.restore()方法載入

如?保存?代碼如下

import tensorflow as tf import numpy as np W = tf.Variable([[1,1,1],[2,2,2]],dtype = tf.float32,name='w') b = tf.Variable([[0,1,2]],dtype = tf.float32,name='b') init = tf.initialize_all_variables() saver = tf.train.Saver() with tf.Session() as sess: sess.run(init) save_path = saver.save(sess,"save/model.ckpt")

載入代碼如下:

import tensorflow as tf import numpy as np W = tf.Variable(tf.truncated_normal(shape=(2,3)),dtype = tf.float32,name='w') b = tf.Variable(tf.truncated_normal(shape=(1,3)),dtype = tf.float32,name='b') saver = tf.train.Saver() with tf.Session() as sess: saver.restore(sess,"save/model.ckpt")

這種方法不方便的在于，在使用模型的時(shí)候，必須把模型的結(jié)構(gòu)重新定義一遍，然后載入對(duì)應(yīng)名字的變量的值。但是很多時(shí)候我們都更希望能夠讀取一個(gè)文件然后就直接使用模型，而不是還要把模型重新定義一遍。所以就需要使用另一種方法。

二、不需重新定義網(wǎng)絡(luò)結(jié)構(gòu)的方法

tf.train.import_meta_graph

import_meta_graph(meta_graph_or_file,clear_devices=False,import_scope=None,**kwargs )

這個(gè)方法可以從文件中將保存的graph的所有節(jié)點(diǎn)加載到當(dāng)前的default graph中，并返回一個(gè)saver。也就是說(shuō)，我們?cè)诒４娴臅r(shí)候，除了將變量的值保存下來(lái)，其實(shí)還有將對(duì)應(yīng)graph中的各種節(jié)點(diǎn)保存下來(lái)，所以模型的結(jié)構(gòu)也同樣被保存下來(lái)了。

比如我們想要保存計(jì)算最后預(yù)測(cè)結(jié)果的y，則應(yīng)該在訓(xùn)練階段將它添加到collection中。具體代碼如下 :

保存

### 定義模型 input_x = tf.placeholder(tf.float32, shape=(None, in_dim), name='input_x') input_y = tf.placeholder(tf.float32, shape=(None, out_dim), name='input_y') w1 = tf.Variable(tf.truncated_normal([in_dim, h1_dim], stddev=0.1), name='w1') b1 = tf.Variable(tf.zeros([h1_dim]), name='b1') w2 = tf.Variable(tf.zeros([h1_dim, out_dim]), name='w2') b2 = tf.Variable(tf.zeros([out_dim]), name='b2') keep_prob = tf.placeholder(tf.float32, name='keep_prob') hidden1 = tf.nn.relu(tf.matmul(self.input_x, w1) + b1) hidden1_drop = tf.nn.dropout(hidden1, self.keep_prob) ### 定義預(yù)測(cè)目標(biāo) y = tf.nn.softmax(tf.matmul(hidden1_drop, w2) + b2) # 創(chuàng)建saver saver = tf.train.Saver(...variables...) # 假如需要保存y，以便在預(yù)測(cè)時(shí)使用 tf.add_to_collection('pred_network', y) sess = tf.Session() for step in xrange(1000000): sess.run(train_op) if step % 1000 == 0: # 保存checkpoint, 同時(shí)也默認(rèn)導(dǎo)出一個(gè)meta_graph # graph名為'my-model-{global_step}.meta'. saver.save(sess, 'my-model', global_step=step)

載入

with tf.Session() as sess:new_saver = tf.train.import_meta_graph('my-save-dir/my-model-10000.meta')new_saver.restore(sess, 'my-save-dir/my-model-10000')# tf.get_collection() 返回一個(gè)list. 但是這里只要第一個(gè)參數(shù)即可 y = tf.get_collection('pred_network')[0] graph = tf.get_default_graph() # 因?yàn)閥中有placeholder，所以sess.run(y)的時(shí)候還需要用實(shí)際待預(yù)測(cè)的樣本以及相應(yīng)的參數(shù)來(lái)填充這些placeholder，而這些需要通過(guò)graph的get_operation_by_name方法來(lái)獲取。 input_x = graph.get_operation_by_name('input_x').outputs[0] keep_prob = graph.get_operation_by_name('keep_prob').outputs[0] # 使用y進(jìn)行預(yù)測(cè) sess.run(y, feed_dict={input_x:...., keep_prob:1.0})

具體示例

save.py

import tensorflow as tf from tensorflow.examples.tutorials.mnist import input_data# 加載數(shù)據(jù)集 mnist = input_data.read_data_sets("data", one_hot=True) # Parameters learning_rate = 0.001 batch_size = 100 display_step = 10 model_path = "save/model.ckpt" # Network Parameters n_hidden_1 = 256 # 1st layer number of features n_hidden_2 = 256 # 2st layer number of features n_input = 784 # MNIST data input (img shape: 28*28) n_classes = 10 # MNIST total classes (0-9 digits) # tf Graph input x = tf.placeholder(tf.float32, [None, n_input], name="input_x") y = tf.placeholder(tf.float32, [None, n_classes], name="input_y") # Store layers weight & bias weights = { 'h1': tf.Variable(tf.random_normal([n_input, n_hidden_1])), 'h2': tf.Variable(tf.random_normal([n_hidden_1, n_hidden_2])), 'out': tf.Variable(tf.random_normal([n_hidden_2, n_classes])) } biases = { 'b1': tf.Variable(tf.random_normal([n_hidden_1])), 'b2': tf.Variable(tf.random_normal([n_hidden_2])), 'out': tf.Variable(tf.random_normal([n_classes])) } # Create model def multilayer_perceptron(x, weights, biases): # layer1 h1 = tf.add(tf.matmul(x, weights['h1']), biases['b1']) h1 = tf.nn.relu(h1) # layer2 h2 = tf.add(tf.matmul(h1, weights['h2']), biases['b2']) h2 = tf.nn.relu(h2) # out out = tf.add(tf.matmul(h2, weights['out']), biases['out']) return out # Construct model logits = multilayer_perceptron(x, weights, biases) pred = tf.nn.softmax(logits) # Define loss and optimizer cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y)) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost) corrcet_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)) accuracy = tf.reduce_mean(tf.cast(corrcet_pred, tf.float32)) # Initializing the variables init = tf.global_variables_initializer() # 保存模型 saver = tf.train.Saver() tf.add_to_collection("pred", pred) tf.add_to_collection('acc', accuracy) with tf.Session() as sess: sess.run(init) step = 0 while step * batch_size < 180000: batch_xs, batch_ys = mnist.train.next_batch(batch_size) loss, _, acc = sess.run([cost, optimizer, accuracy], feed_dict={x: batch_xs, y: batch_ys}) if step % display_step == 0: # step: 1790 loss: 16.9724 acc: 0.95 print("step: ", step, "loss: ", loss, "acc: ", acc) saver.save(sess, save_path=model_path, global_step=step) step += 1 print("Train Finish!")

checkpoint:

model_checkpoint_path: "model.ckpt-1790" all_model_checkpoint_paths: "model.ckpt-1750" all_model_checkpoint_paths: "model.ckpt-1760" all_model_checkpoint_paths: "model.ckpt-1770" all_model_checkpoint_paths: "model.ckpt-1780" all_model_checkpoint_paths: "model.ckpt-1790"

restore.py

import tensorflow as tf from tensorflow.examples.tutorials.mnist import input_data# load mnist data mnist = input_data.read_data_sets("data", one_hot=True) with tf.Session() as sess: new_saver = tf.train.import_meta_graph("save/model.ckpt-1790.meta") new_saver.restore(sess, "save/model.ckpt-1790") # tf.get_collection() 返回一個(gè)list. 但是這里只要第一個(gè)參數(shù)即可 pred = tf.get_collection("pred")[0] acc = tf.get_collection("acc")[0] # 因?yàn)?pred, acc 中有 placeholder，所以 sess.run(acc)的時(shí)候還需要用實(shí)際待預(yù)測(cè)的樣本以及相應(yīng)的參數(shù)來(lái)填充這些placeholder， # 而這些需要通過(guò)graph的get_operation_by_name方法來(lái)獲取。 graph = tf.get_default_graph() x = graph.get_operation_by_name("input_x").outputs[0] y = graph.get_operation_by_name("input_y").outputs[0] test_xs = mnist.test.images test_ys = mnist.test.labels #test set acc: [0.91820002] print("test set acc: ", sess.run([acc], feed_dict={ x: test_xs, y: test_ys }))
原文：https://blog.csdn.net/u011026329/article/details/79190347

總結(jié)

以上是生活随笔為你收集整理的TensorFlow模型保存和加载方法的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。

上一篇：【BZOJ 1877】 [SDOI200
下一篇： Python全栈 MongoDB 数据库