start
首先將本文涉及到的代碼、預訓練模型和數據集上傳，大家可以自行下載：
VGG16代碼鏈接：
https://pan.baidu.com/s/1Xy5H3t9SVnQM2OMorH4pmQ
提取碼：zju1
預訓練模型VGG16.npy鏈接：
https://pan.baidu.com/s/1-HQL4Ixkm8G2j01y8sstFA
提取碼：o9g2
數據集鏈接：
https://pan.baidu.com/s/1z9Y5L1B10huqGF-nBu0mug
提取碼：fqnb
01文件目錄介紹
“
本文介紹使用tensorflow實現VGG16網絡模型並進行圖像分類，主文件夾VGG16_Tensorflow下目錄如下：
”
其中dataset文件夾下爲數據集文件，model文件夾爲訓練模型的儲存目錄，test文件夾爲測試集圖像的存放位置。
主文件夾下包含4個py文件，分別實現不同的功能：
create_tfrecords.py 生成tfrecords數據腳本
VGG16.py 網絡結構定義文件
train.py 訓練腳本
test.py 測試腳本
dataset文件夾包含子目錄data：
data目錄下式五個類別的圖片
test文件夾下存放的是測試集圖像，如下：
02製作tfrecord數據文件
“
首先，製作數據集並按照如下格式保存到dataset/data/文件夾下
”
然後，生成train.tfrecords文件
按照如上目錄要求製作完數據集後，在主目錄下直接運行python create_tfrecords.py即可在主目錄下生成train.tfrecords文件。以下是create_tfrecords.py代碼：
#coding=utf-8
import os
import tensorflow as tf
from PIL import Image
import sys
def creat_tf(imgpath):
cwd = os.getcwd()
classes = os.listdir(cwd + imgpath)
writer = tf.python_io.TFRecordWriter("train.tfrecords")
for index, name in enumerate(classes):
class_path = cwd + imgpath + name + "/"
print(class_path)
if os.path.isdir(class_path):
for img_name in os.listdir(class_path):
img_path = class_path + img_name
img = Image.open(img_path)
img = img.resize((224, 224))
img_raw = img.tobytes() #將圖片轉化爲原生bytes
example = tf.train.Example(features=tf.train.Features(feature={
'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[int(name)])),
'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
}))
writer.write(example.SerializeToString()) #序列化爲字符串
print(img_name)
writer.close()
def read_example():
#簡單的讀取例子：
for serialized_example in tf.python_io.tf_record_iterator("train.tfrecords"):
example = tf.train.Example()
example.ParseFromString(serialized_example)
#image = example.features.feature['img_raw'].bytes_list.value
label = example.features.feature['label'].int64_list.value
# 可以做一些預處理之類的
print(label)
if __name__ == '__main__':
imgpath = './dataset/data/'# 訓練集圖像的儲存目錄
creat_tf(imgpath)
03模型訓練
“
製作完數據後便可進行模型的訓練工作，直接運行python train.py即可。
”
模型訓練過程中會調用預訓練模型文件vgg16.npy，可以自行從下面的鏈接中下載：
鏈接：
https://pan.baidu.com/s/1-HQL4Ixkm8G2j01y8sstFA
提取碼：o9g2
#coding=utf-8
import tensorflow as tf
import numpy as np
import pdb
from datetime import datetime
from VGG16 import *
batch_size = 64#批大小
lr = 0.00001#學習率
n_cls = 17
max_steps = 10000#訓練次數
def read_and_decode(filename):
#根據文件名生成一個隊列
filename_queue = tf.train.string_input_producer([filename])
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue) #返回文件名和文件
features = tf.parse_single_example(serialized_example,
features={
'label': tf.FixedLenFeature([], tf.int64),
'img_raw' : tf.FixedLenFeature([], tf.string),
})
img = tf.decode_raw(features['img_raw'], tf.uint8)
img = tf.reshape(img, [224, 224, 3])
# 轉換爲float32類型，並做歸一化處理
img = tf.cast(img, tf.float32)# * (1. / 255)
label = tf.cast(features['label'], tf.int64)
#print 'images的樣子是：', img
#print 'label的樣子是：', label
#pdb.set_trace()
return img, label
def train():
x = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3], name='input')
y = tf.placeholder(dtype=tf.float32, shape=[None, n_cls], name='label')
keep_prob = tf.placeholder(tf.float32)
output = VGG16(x, keep_prob, n_cls)
#probs = tf.nn.softmax(output)
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=output, labels=y))
#train_step = tf.train.AdamOptimizer(learning_rate=0.1).minimize(loss)
train_step = tf.train.GradientDescentOptimizer(learning_rate=lr).minimize(loss)
accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(output,1), tf.argmax(y, 1)), tf.float32))
images, labels = read_and_decode('./train.tfrecords')
img_batch, label_batch = tf.train.shuffle_batch([images, labels],
batch_size=batch_size,
capacity=392,
min_after_dequeue=200)
label_batch = tf.one_hot(label_batch, n_cls, 1, 0)
init = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
for i in range(max_steps):
batch_x, batch_y = sess.run([img_batch, label_batch])
# print batch_x, batch_x.shape
# print batch_y
# pdb.set_trace()
_, loss_val = sess.run([train_step, loss], feed_dict={x:batch_x, y:batch_y, keep_prob:0.8})
if i%10 == 0:
train_arr = accuracy.eval(feed_dict={x:batch_x, y: batch_y, keep_prob: 1.0})
print("%s: Step [%d] Loss : %f, training accuracy : %g" % (datetime.now(), i, loss_val, train_arr))
if (i + 1) == max_steps:
#checkpoint_path = os.path.join(FLAGS.train_dir, './model/model.ckpt')
saver.save(sess, './model/model.ckpt', global_step=i)
coord.request_stop()
coord.join(threads)
#saver.save(sess, 'model/model.ckpt')
if __name__ == '__main__':
train()
模型訓練結束後會在model文件夾下生成模型文件，以我的數據集爲例，訓練10000步後生成如下文件：
04模型測試
♡♡♡
模型訓練完畢後運行python test.py即可完成模型測試工作，test.py文件代碼如下：
#coding=utf-8
import tensorflow as tf
import numpy as np
import pdb
from datetime import datetime
from VGG16 import *
import cv2
import os
def test(path):
x = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3], name='input')
keep_prob = tf.placeholder(tf.float32)
output = VGG16(x, keep_prob, 17)
score = tf.nn.softmax(output)
f_cls = tf.argmax(score, 1)
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, './model/model.ckpt-9999')#調用模型的名稱和路徑
for i in os.listdir(path):
imgpath = os.path.join(path, i)
im = cv2.imread(imgpath)
im = cv2.resize(im, (224 , 224))# * (1. / 255)
im = np.expand_dims(im, axis=0)
#pred = sess.run(f_cls, feed_dict={x:im, keep_prob:1.0})
pred, _score = sess.run([f_cls, score], feed_dict={x:im, keep_prob:1.0})
prob = round(np.max(_score), 4)
#print "{} flowers class is: {}".format(i, pred)
print("{} flowers class is: {}, score: {}".format(i, int(pred), prob))
sess.close()
if __name__ == '__main__':
path = './test'#測試集圖像路徑
test(path)
以我當前的數據集爲例進行測試，測試結果如下：
大家在訓練過程中如有疑問，歡迎後臺留言討論！
參考：
https://github.com/LiMingda92/VGG16_TF
