tensorflow迁移学习

最新推荐文章于 2023-08-22 00:00:21 发布

joyjoy01

最新推荐文章于 2023-08-22 00:00:21 发布

阅读量313

点赞数

CC 4.0 BY-SA版权

分类专栏： tensorflow

本文链接：https://blog.youkuaiyun.com/joyjoy01/article/details/80361880

tensorflow 专栏收录该内容

4 篇文章

订阅专栏

本文介绍了一种基于TensorFlow的手写数字模型迁移方法。通过冻结预训练模型的部分层，并添加新的全连接层来适应新的任务。该过程包括加载MNIST数据集、定义计算图、进行迁移学习并最终保存迁移后的模型。

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

模型是我自己保存的手写字体模型

#!/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 np

def 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]})