tensorflow實戰—VGGNET16–遷移學習訓練自己的數據集

tensorflow實戰—VGGNET16–遷移學習訓練自己的數據集

本文參考----https://blog.csdn.net/wangds000/article/details/99670777
先介紹一下自己的電腦配置環境吧。這次試驗用cpu跑的，電腦內存24g，沒用gpu顯卡就不介紹了。64位操作系統，處理器i7.
數據集採用的是ISIC-2018任務三中皮膚鏡數據集。
在參考的博主遷移學習vggnet16的網絡模型失效了，我在網上找了個，存放在百度雲了，有需要的朋友可以下載。
鏈接：https://pan.baidu.com/s/1ojhtxjX-XVb2ONBAp-ozQA
提取碼：ykim

數據介紹

下面的是自己的數據集存放位置。(0 AKIEC 1 BCC 2 BKL 3 DF 4 MEL 5 NV 6 VASC)七個病種數據
-DATA
----0
----1
----2
----3
----4
----5
----6
-logs
-model
-test
creat_tfrecords.py
test.py
train.py
VGG16.py
vgg16.npy(vgg16網絡模型)

數據轉換

將七類數據集轉化爲tfrecord格式，imgpath爲你訓練數據的路徑，代碼如下

import os
import tensorflow as tf
from PIL import Image
import sys
import matplotlib.pyplot as plt

def creat_tf(imgpath):
 
    cwd = os.getcwd()
    classes = os.listdir(cwd + imgpath)
    
    # 此處定義tfrecords文件存放
    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()             
                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 = 'DATA/'
    creat_tf(imgpath)

加載vggnet16遷移學習模型

可能存在tensorflow版本問題，tf.reset_default_graph()模塊引用不正確，讀者可以百度一下自己tensorflow的版本信息，選用正確的模塊書寫格式。

import tensorflow as tf
import numpy as np 
tf.reset_default_graph()
 
# 加載預訓練模型
data_dict = np.load('vgg16.npy', encoding='latin1').item()

# 打印每層信息
def print_layer(t):
    print(t.op.name, ' ', t.get_shape().as_list(), '\n')
 
# 定義卷積層
"""
此處權重初始化定義了3種方式：
    1.預訓練模型參數
    2.截尾正態，參考書上採用該方式
    3.xavier，網上blog有采用該方式
通過參數finetrun和xavier控制選擇哪種方式，有興趣的可以都試試
"""
def conv(x, d_out, name, fineturn=False, xavier=False):
    d_in = x.get_shape()[-1].value
    with tf.name_scope(name) as scope:
        # Fine-tuning 
        if fineturn:
            kernel = tf.constant(data_dict[name][0], name="weights")
            bias = tf.constant(data_dict[name][1], name="bias")
            #print("fineturn")
        elif not xavier:
            kernel = tf.Variable(tf.truncated_normal([3, 3, d_in, d_out], stddev=0.1), name='weights')
            bias = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[d_out]),
                                                trainable=True, 
                                                name='bias')
            #print("truncated_normal")
        else:
            kernel = tf.get_variable(scope+'weights', shape=[3, 3, d_in, d_out], 
                                                dtype=tf.float32,
                                                initializer=tf.contrib.layers.xavier_initializer_conv2d())
            bias = tf.Variable(tf.constant(0.0, dtype=tf.float32, shape=[d_out]),
                                                trainable=True, 
                                                name='bias')
            #print("xavier")
        conv = tf.nn.conv2d(x, kernel,[1, 1, 1, 1], padding='SAME')
        activation = tf.nn.relu(conv + bias, name=scope)
        #print_layer(activation)
        return activation
 
# 最大池化層
def maxpool(x, name):
    activation = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID', name=name) 
    #print_layer(activation)
    return activation
 
# 定義全連接層
"""
此處權重初始化定義了3種方式：
    1.預訓練模型參數
    2.截尾正態，參考書上採用該方式
    3.xavier，網上blog有采用該方式
通過參數finetrun和xavier控制選擇哪種方式，有興趣的可以都試試
"""
def fc(x, n_out, name, fineturn=False, xavier=False):
    n_in = x.get_shape()[-1].value
    with tf.name_scope(name) as scope:
        if fineturn:
            weight = tf.constant(data_dict[name][0], name="weights")
            bias = tf.constant(data_dict[name][1], name="bias")
            #print("fineturn")
        elif not xavier:
            weight = tf.Variable(tf.truncated_normal([n_in, n_out], stddev=0.01), name='weights')
            bias = tf.Variable(tf.constant(0.1, dtype=tf.float32, shape=[n_out]), 
                                                trainable=True, 
                                                name='bias')
            #print("truncated_normal")
        else:
            weight = tf.get_variable(scope+'weights', shape=[n_in, n_out], 
                                                dtype=tf.float32,
                                                initializer=tf.contrib.layers.xavier_initializer_conv2d())
            bias = tf.Variable(tf.constant(0.1, dtype=tf.float32, shape=[n_out]), 
                                                trainable=True, 
                                                name='bias')
            #print("xavier")
        # 全連接層可以使用relu_layer函數比較方便，不用像卷積層使用relu函數
        activation = tf.nn.relu_layer(x, weight, bias, name=name)
        #print_layer(activation)
        return activation
 
def VGG_16(images, _dropout, n_cls):
    """
    此處權重初始化方式採用的是：
        卷積層使用預訓練模型中的參數
        全連接層使用xavier類型初始化
    """
    conv1_1 = conv(images, 64, 'conv1_1', fineturn=True)
    conv1_2 = conv(conv1_1, 64, 'conv1_2', fineturn=True)
    pool1   = maxpool(conv1_2, 'pool1')
 
    conv2_1 = conv(pool1, 128, 'conv2_1', fineturn=True)
    conv2_2 = conv(conv2_1, 128, 'conv2_2', fineturn=True)
    pool2   = maxpool(conv2_2, 'pool2')
 
    conv3_1 = conv(pool2, 256, 'conv3_1', fineturn=True)
    conv3_2 = conv(conv3_1, 256, 'conv3_2', fineturn=True)
    conv3_3 = conv(conv3_2, 256, 'conv3_3', fineturn=True)
    pool3   = maxpool(conv3_3, 'pool3')
 
    conv4_1 = conv(pool3, 512, 'conv4_1', fineturn=True)
    conv4_2 = conv(conv4_1, 512, 'conv4_2', fineturn=True)
    conv4_3 = conv(conv4_2, 512, 'conv4_3', fineturn=True)
    pool4   = maxpool(conv4_3, 'pool4')
 
    conv5_1 = conv(pool4, 512, 'conv5_1', fineturn=True)
    conv5_2 = conv(conv5_1, 512, 'conv5_2', fineturn=True)
    conv5_3 = conv(conv5_2, 512, 'conv5_3', fineturn=True)
    pool5   = maxpool(conv5_3, 'pool5')
 
    '''
    因爲訓練自己的數據，全連接層最好不要使用預訓練參數
    '''
    flatten  = tf.reshape(pool5, [-1, 7*7*512])
    fc6      = fc(flatten, 4096, 'fc6', xavier=True)
    dropout1 = tf.nn.dropout(fc6, _dropout)
 
    fc7      = fc(dropout1, 4096, 'fc7', xavier=True)
    dropout2 = tf.nn.dropout(fc7, _dropout)
    
    fc8      = fc(dropout2, n_cls, 'fc8', xavier=True)
 
    return fc8

訓練

博主訓練次數採用15000次，我實驗的時候沒有改，發生keyboardinterrupt錯誤，百度了下可能是線程中斷問題，我猜想可能與我訓練次數太多，cpu處理器處理失敗，線程中斷，所以我把訓練次數改成了400

import tensorflow as tf 
#import numpy as np 
#import pdb
from datetime import datetime
#from VGG16 import *
import VGG16
#tf.reset_default_graph()

batch_size = 24
lr = 0.0001
n_cls = 29 #訓練時根據自己的類別數更改
max_steps = 400
 
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)
    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.VGG_16(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)
    # log彙總記錄
    summary_op = tf.summary.merge_all()

    init = tf.global_variables_initializer()
    #saver = tf.train.Saver()
    saver = tf.train.Saver(max_to_keep=3)
    max_acc = 0
    with tf.Session() as sess:
        sess.run(init)
        train_writer = tf.summary.FileWriter('logs',sess.graph)
        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)
                summary_str = sess.run(summary_op)
                train_writer.add_summary(summary_str,i)
            # 保存最近的3個模型參數
            if train_arr > max_acc:
                saver.save(sess, 'model/model.ckpt', global_step=i + 1)
        coord.request_stop()
        coord.join(threads)
        #saver.save(sess, 'model/model.ckpt')
		
 
if __name__ == '__main__':
    train()

測試

測試需要用到opencv，此時需要安裝opencv庫
命令：pip install opencv-python
注意：
saver.restore(sess, ‘model/model.ckpt-442’)
這個加載你的最後一次訓練結果的模型（我是訓練的400次，出來442我也不清楚了）

import tensorflow as tf 
import numpy as np 
import pdb
from datetime import datetime
from VGG16 import *
import cv2
import os
#import matplotlib.pyplot as plt
 
def test(path):
 
    x = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3], name='input')
    keep_prob = tf.placeholder(tf.float32)
    # 注意更改自己的類別數，此處輸出爲29類
    output = VGG_16(x, keep_prob, 7)
    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-442')
    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)
        # 測試時，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("{} rubbing class is: {}, score: {}".format(i, int(pred), prob))
        # plt.imshow(im)
        # plt.imshow(im1)
        # plt.title(u'預測值:%i' % pred)
        # plt.show() 
    sess.close()
    
if __name__ == '__main__':
    # 測試圖片保存在文件夾中了，圖片前面數字爲所屬類別
    path = 'C:/Users/Administrator/Desktop/vggnet19/test/'
    test(path)

結果

訓練結果截圖：

測試結果截圖：