模型是我自己保存的手写字体模型
#!/usr/bin/python
# -*- coding:utf-8 -*- import input_data as input import tensorflow as tf from tensorflow.python.framework import graph_util from tensorflow.python.platform import gfile import numpy as npdef transfer():
mnist = input.read_data_sets("/home/myjob/Downloads/Mnist/", one_hot=True) model_path_pb = "/home/myjob/Downloads/Mnist/" y_ = tf.placeholder("float32", [None, 10], name='input_y') with tf.Session() as sess:#读取需要迁移的模型
with gfile.FastGFile(model_path_pb + 'model.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) sess.graph.as_default() tf.import_graph_def(graph_def, name='') # 导入计算图 input_x = sess.graph.get_tensor_by_name('input_x:0') h_fc = sess.graph.get_tensor_by_name('fc1:0') # h_fc层以上的层停止梯度传递，相当于keras中的freeze hfc_sg = tf.stop_gradient(h_fc) #添加新的层 w_fc1 = weight_variable([1024, 1024]) b_fc1 = bias_Variable([1024]) h_fc1 = tf.nn.relu(tf.matmul(hfc_sg, w_fc1) + b_fc1, name="fc2") w_fc2 = weight_variable([1024, 10]) b_fc2 = bias_Variable([10]) y_conv2d = tf.nn.softmax(tf.matmul(h_fc1, w_fc2) + b_fc2, name="outt") #交叉熵损失函数 cross_entropy = -tf.reduce_sum(y_ * tf.log(y_conv2d)) train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy) correct_prediction = tf.equal(tf.argmax(y_conv2d, 1), tf.argmax(y_, 1)) accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float32")) init = tf.global_variables_initializer() with tf.Session() as sess: sess.run(init) for i in range(101): batch = mnist.train.next_batch(64) if i % 2 == 0: train_accuracy = accuracy.eval(feed_dict={ input_x: batch[0], y_: batch[1]}) print "step %d, training accuracy %g" % (i, train_accuracy) #将迁移后的模型保存 constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['outt']) with tf.gfile.FastGFile(model_path_pb + 'model_transform.pb', mode='wb') as f: f.write(constant_graph.SerializeToString()) train_step.run(feed_dict={input_x:batch[0], y_: batch[1]})