当前位置：首页 > news >正文

如何用Tensorflow训练模型成pb文件和和如何加载已经训练好的模型文件

news 来源：原创 2024/5/6 7:54:09

这篇薄荷主要是讲了如何用tensorflow去训练好一个模型，然后生成相应的pb文件。最后会将如何重新加载这个pb文件。
首先先放出PO主的github：

https://github.com/ppplinday/tensorflow-vgg16-train-and-test

其中的pitcute文件是狗和猫的图片分别15张一共30（别吐槽，只是为了练手学习的233333）， train那个就是训练的文件，test这个就是测试的文件。
接着PO主会慢慢讲解相应的步骤。
！！！ps：由于PO主也是新手，所以难免会出现一点（很多）小错误，希望大婶看了能够提出来让PO主好好学习233333。

train
首先说一下train。一开始当然是读图片啦。

def read_img(path):
    cate   = [path + x for x in os.listdir(path) if os.path.isdir(path + x)] imgs = [] labels = [] for idx, folder in enumerate(cate): for im in glob.glob(folder + '/*.jpg'): print('reading the image: %s' % (im)) img = io.imread(im) img = transform.resize(img, (w, h, c)) imgs.append(img) labels.append(idx) return np.asarray(imgs, np.float32), np.asarray(labels, np.int32) data, label = read_img(path)

用io.imread来读取每一张图片，然后resize成vgg的输入的大小（224，224，3），最后分别放入了data和label中。

num_example = data.shape[0]
arr = np.arange(num_example)
np.random.shuffle(arr) data = data[arr] label = label[arr]

这里是把图片的顺序打乱，先生成一个等差数列，然后打乱，最后赋值回原来的data和label

ratio = 0.8
s = np.int(num_example * ratio) x_train = data[:s] y_train = label[:s] x_val = data[s:] y_val = label[s:]

全部的数据中百分之80的用来train，剩下20的用来test（虽然一共才30张图片。。。。。）

def build_network(height, width, channel):
    x = tf.placeholder(tf.float32, shape=[None, height, width, channel], name='input') y = tf.placeholder(tf.int64, shape=[None, 2], name='labels_placeholder')

开始build相应的vgg model，这一步不难，但是每一层最好都给上相应的name。上面的x和y是相应的输入和相应的标签。

    finaloutput = tf.nn.softmax(output_fc8, name="softmax")

    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=finaloutput, labels=y)) optimize = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(cost) prediction_labels = tf.argmax(finaloutput, axis=1, name="output") read_labels = y correct_prediction = tf.equal(prediction_labels, read_labels) accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) correct_times_in_batch = tf.reduce_sum(tf.cast(correct_prediction, tf.int32)) return dict( x=x, y=y, optimize=optimize, correct_prediction=correct_prediction, correct_times_in_batch=correct_times_in_batch, cost=cost, )

在build的最后，是需要进行误差计算。finaloutput是最后的输出，cost是计算误差，optimize是定义训练时候安什么方式，也注意一下最后的return。

接着是训练过程。

def train_network(graph, batch_size, num_epochs, pb_file_path):
    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        epoch_delta = 2
        for epoch_index in range(num_epochs):
            for i in range(12): sess.run([graph['optimize']], feed_dict={ graph['x']: np.reshape(x_train[i], (1, 224, 224, 3)), graph['y']: ([[1, 0]] if y_train[i] == 0 else [[0, 1]]) })

其实训练的代码就这些，定好了batchsize和numepoch进行训练。下面的代码主要是为了看每几次相应的正确率。

            constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["output"])
            with tf.gfile.FastGFile(pb_file_path, mode='wb') as f:
f.write(constant_graph.SerializeToString())

这两句是重要的代码，用来把训练好的模型保存为pb文件。运行完之后就会发现应该的文件夹多出了一个pb文件。

test

def recognize(jpg_path, pb_file_path):
    with tf.Graph().as_default(): output_graph_def = tf.GraphDef() with open(pb_file_path, "rb") as f: output_graph_def.ParseFromString(f.read()) _ = tf.import_graph_def(output_graph_def, name="")

打开相应的pb文件。

            img = io.imread(jpg_path)
            img = transform.resize(img, (224, 224, 3)) img_out_softmax = sess.run(out_softmax, feed_dict={input_x:np.reshape(img, [-1, 224, 224, 3])})

读取图片文件，resize之后放入模型的输入位置，之后img_out_softmax就是相应输出的结果。

这大概就是整个流程。目的是为了练练手，PO主应该有挺多小错误，希望大家能够提出来让PO主好好学习哈哈哈！！！

最后放出整个的train和test的代码：
train

from PIL import Image
import numpy as np
import matplotlib.pyplot as plt import matplotlib.image as mpimg import tensorflow as tf import os import glob from skimage import io, transform from tensorflow.python.framework import graph_util import collections path = '/home/zhoupeilin/vgg16/picture/' w = 224 h = 224 c = 3 def read_img(path): cate = [path + x for x in os.listdir(path) if os.path.isdir(path + x)] imgs = [] labels = [] for idx, folder in enumerate(cate): for im in glob.glob(folder + '/*.jpg'): print('reading the image: %s' % (im)) img = io.imread(im) img = transform.resize(img, (w, h, c)) imgs.append(img) labels.append(idx) return np.asarray(imgs, np.float32), np.asarray(labels, np.int32) data, label = read_img(path) num_example = data.shape[0] arr = np.arange(num_example) np.random.shuffle(arr) data = data[arr] label = label[arr] ratio = 0.8 s = np.int(num_example * ratio) x_train = data[:s] y_train = label[:s] x_val = data[s:] y_val = label[s:] def build_network(height, width, channel): x = tf.placeholder(tf.float32, shape=[None, height, width, channel], name='input') y = tf.placeholder(tf.int64, shape=[None, 2], name='labels_placeholder') def weight_variable(shape, name="weights"): initial = tf.truncated_normal(shape, dtype=tf.float32, stddev=0.1) return tf.Variable(initial, name=name) def bias_variable(shape, name="biases"): initial = tf.constant(0.1, dtype=tf.float32, shape=shape) return tf.Variable(initial, name=name) def conv2d(input, w): return tf.nn.conv2d(input, w, [1, 1, 1, 1], padding='SAME') def pool_max(input): return tf.nn.max_pool(input, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool1') def fc(input, w, b): return tf.matmul(input, w) + b # conv1 with tf.name_scope('conv1_1') as scope: kernel = weight_variable([3, 3, 3, 64]) biases = bias_variable([64]) output_conv1_1 = tf.nn.relu(conv2d(x, kernel) + biases, name=scope) with tf.name_scope('conv1_2') as scope: kernel = weight_variable([3, 3, 64, 64]) biases = bias_variable([64]) output_conv1_2 = tf.nn.relu(conv2d(output_conv1_1, kernel) + biases, name=scope) pool1 = pool_max(output_conv1_2) # conv2 with tf.name_scope('conv2_1') as scope: kernel = weight_variable([3, 3, 64, 128]) biases = bias_variable([128]) output_conv2_1 = tf.nn.relu(conv2d(pool1, kernel) + biases, name=scope) with tf.name_scope('conv2_2') as scope: kernel = weight_variable([3, 3, 128, 128]) biases = bias_variable([128]) output_conv2_2 = tf.nn.relu(conv2d(output_conv2_1, kernel) + biases, name=scope) pool2 = pool_max(output_conv2_2) # conv3 with tf.name_scope('conv3_1') as scope: kernel = weight_variable([3, 3, 128, 256]) biases = bias_variable([256]) output_conv3_1 = tf.nn.relu(conv2d(pool2, kernel) + biases, name=scope) with tf.name_scope('conv3_2') as scope: kernel = weight_variable([3, 3, 256, 256]) biases = bias_variable([256]) output_conv3_2 = tf.nn.relu(conv2d(output_conv3_1, kernel) + biases, name=scope) with tf.name_scope('conv3_3') as scope: kernel = weight_variable([3, 3, 256, 256]) biases = bias_variable([256]) output_conv3_3 = tf.nn.relu(conv2d(output_conv3_2, kernel) + biases, name=scope) pool3 = pool_max(output_conv3_3) # conv4 with tf.name_scope('conv4_1') as scope: kernel = weight_variable([3, 3, 256, 512]) biases = bias_variable([512]) output_conv4_1 = tf.nn.relu(conv2d(pool3, kernel) + biases, name=scope) with tf.name_scope('conv4_2') as scope: kernel = weight_variable([3, 3, 512, 512]) biases = bias_variable([512]) output_conv4_2 = tf.nn.relu(conv2d(output_conv4_1, kernel) + biases, name=scope) with tf.name_scope('conv4_3') as scope: kernel = weight_variable([3, 3, 512, 512]) biases = bias_variable([512]) output_conv4_3 = tf.nn.relu(conv2d(output_conv4_2, kernel) + biases, name=scope) pool4 = pool_max(output_conv4_3) # conv5 with tf.name_scope('conv5_1') as scope: kernel = weight_variable([3, 3, 512, 512]) biases = bias_variable([512]) output_conv5_1 = tf.nn.relu(conv2d(pool4, kernel) + biases, name=scope) with tf.name_scope('conv5_2') as scope: kernel = weight_variable([3, 3, 512, 512]) biases = bias_variable([512]) output_conv5_2 = tf.nn.relu(conv2d(output_conv5_1, kernel) + biases, name=scope) with tf.name_scope('conv5_3') as scope: kernel = weight_variable([3, 3, 512, 512]) biases = bias_variable([512]) output_conv5_3 = tf.nn.relu(conv2d(output_conv5_2, kernel) + biases, name=scope) pool5 = pool_max(output_conv5_3) #fc6 with tf.name_scope('fc6') as scope: shape = int(np.prod(pool5.get_shape()[1:])) kernel = weight_variable([shape, 4096]) biases = bias_variable([4096]) pool5_flat = tf.reshape(pool5, [-1, shape]) output_fc6 = tf.nn.relu(fc(pool5_flat, kernel, biases), name=scope) #fc7 with tf.name_scope('fc7') as scope: kernel = weight_variable([4096, 4096]) biases = bias_variable([4096]) output_fc7 = tf.nn.relu(fc(output_fc6, kernel, biases), name=scope) #fc8 with tf.name_scope('fc8') as scope: kernel = weight_variable([4096, 2]) biases = bias_variable([2]) output_fc8 = tf.nn.relu(fc(output_fc7, kernel, biases), name=scope) finaloutput = tf.nn.softmax(output_fc8, name="softmax") cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=finaloutput, labels=y)) optimize = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(cost) prediction_labels = tf.argmax(finaloutput, axis=1, name="output") read_labels = y correct_prediction = tf.equal(prediction_labels, read_labels) accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) correct_times_in_batch = tf.reduce_sum(tf.cast(correct_prediction, tf.int32)) return dict( x=x, y=y, optimize=optimize, correct_prediction=correct_prediction, correct_times_in_batch=correct_times_in_batch, cost=cost, ) def train_network(graph, batch_size, num_epochs, pb_file_path): init = tf.global_variables_initializer() with tf.Session() as sess: sess.run(init) epoch_delta = 2 for epoch_index in range(num_epochs): for i in range(12): sess.run([graph['optimize']], feed_dict={ graph['x']: np.reshape(x_train[i], (1, 224,