迁移学习及应用

猫狗大战
强烈建议GPU环境做深度学习，不然CPU卡到你怀疑人生！！！

源代码：

import numpy as np
import tensorflow as tf
import os
import tensorflow as tf
from time import time
from vgg_preprocess import preprocess_for_train
tf.reset_default_graph() 

class vgg16:
    def __init__(self,imgs):
        self.parameters = []
        self.imgs = imgs
        self.convlayers()
        self.fc_layers()
        self.probs = tf.nn.softmax(self.fc8)
        
    def saver(self):
        return tf.train.Saver()
    


    #定义池化层
    def maxpool(self,name,input_data):
        out = tf.nn.max_pool(input_data,[1,2,2,1],[1,2,2,1],padding="SAME",name=name)
        return out

    #定义卷积层
    def conv(self,name,input_data,out_channel,trainable=False):
        in_channel = input_data.get_shape()[-1]
        with tf.variable_scope(name):
            kernel = tf.get_variable("kernel",[3,3,in_channel,out_channel],dtype=tf.float32,trainable=False)
            Biases = tf.get_variable("Biases",[out_channel],dtype=tf.float32,trainable=False)
            conv_res = tf.nn.conv2d(input_data,kernel,[1,1,1,1],padding="SAME")
            res = tf.nn.bias_add(conv_res,Biases)
            out = tf.nn.relu(res,name=name)
        self.parameters += [kernel,Biases]
        return out

    #定义全连接层
    def fc(self,name,input_data,out_channel,trainable=True):
        shape = input_data.get_shape().as_list()
        if len(shape) == 4:
            size = shape[-1]*shape[-2]*shape[-3]
        else:
            size = shape[1]
        #对数据进行展开操作
        input_data_flat = tf.reshape(input_data,[-1,size])
        with tf.variable_scope(name):
            weights = tf.get_variable(name="weights",shape=[size,out_channel],dtype=tf.float32,trainable=trainable)
            biases = tf.get_variable(name="biases",shape=[out_channel],dtype=tf.float32,trainable=trainable)
            res = tf.matmul(input_data_flat,weights)
            out = tf.nn.relu(tf.nn.bias_add(res,biases))
        self.parameters += [weights,biases]
        return out
    
    def convlayers(self):
        #conv1
        self.conv1_1 = self.conv("conv1re_1",self.imgs,64,trainable=False)
        self.conv1_2 = self.conv("conv1_2",self.conv1_1,64,trainable=False)
        self.pool1 = self.maxpool("poolre1",self.conv1_2)
        
        #conv2
        self.conv2_1 = self.conv("conv2_1",self.pool1,128,trainable=False)
        self.conv2_2 = self.conv("convwe2_2",self.conv2_1,128,trainable=False)
        self.pool2 = self.maxpool("pool2",self.conv2_2)
        
        #conv3
        self.conv3_1 = self.conv("conv3_1",self.pool2,256,trainable=False)
        self.conv3_2 = self.conv("convrwe3_2",self.conv3_1,256,trainable=False)
        self.conv3_3 = self.conv("convrwe3_3",self.conv3_2,256,trainable=False)
        self.pool3 = self.maxpool("poolre3",self.conv3_3)
        
        #conv4
        self.conv4_1 = self.conv("conv4_1",self.pool3,512,trainable=False)
        self.conv4_2 = self.conv("convrwe4_2",self.conv4_1,512,trainable=False)
        self.conv4_3 = self.conv("conv4rwe_3",self.conv4_2,512,trainable=False)
        self.pool4 = self.maxpool("pool4",self.conv4_3)
        
        #conv5
        self.conv5_1 = self.conv("conv5_1",self.pool4,512,trainable=False)
        self.conv5_2 = self.conv("convrwe5_2",self.conv5_1,512,trainable=False)
        self.conv5_3 = self.conv("conv5_3",self.conv5_2,512,trainable=False)
        self.pool5 = self.maxpool("poorwel5",self.conv5_3)
        
    
    def fc_layers(self):
        self.fc6 = self.fc("fc1",self.pool5,4096,trainable=False)
        self.fc7 = self.fc("fc2",self.fc6,4096,trainable=False)
        self.fc8 = self.fc("fc3",self.fc7,2,trainable=True) #2分类
        
        
    #载入权重
    def load_weights(self,weight_file,sess):
        weights = np.load(weight_file)
        keys = sorted(weights.keys())
        for i,k in enumerate(keys):
            if i not in [30,31]:
                sess.run(self.parameters[i].assign(weights[k]))
        print("----------Weights loaded----------")



class U:
    #数据输入
    def get_file(self,file_dir):
        images = []
        temp =[]
        for root,sub_folders,files in os.walk(file_dir):
            for name in files:
                images.append(os.path.join(root,name))
            for name in sub_folders:
                temp.append(os.path.join(root,name))
                labels = []
        for one_folder in temp:
            n_img = len(os.listdir(one_folder))
            letter = one_folder.split("/")[-1]
            if letter == "cat":
                labels = np.append(labels,n_img*[0])
            else:
                labels = np.append(labels,n_img*[1])

        #shuffle
        temp = np.array([images,labels])
        temp = temp.transpose()
        np.random.shuffle(temp)
        image_list = list(temp[:,0])
        label_list = list(temp[:,1])
        label_list = [int(float(i)) for i in label_list]

        return image_list,label_list

    img_width = 224
    img_height = 224

    def get_batch(self,image_list,label_list,img_width,img_height,batch_size,capacity):
        image = tf.cast(image_list,tf.string)
        label = tf.cast(label_list,tf.int32)
        input_queue = tf.train.slice_input_producer([image,label])
        label = input_queue[1]
        image_contents = tf.read_file(input_queue[0])

        image = tf.image.decode_jpeg(image_contents,channels=3)
        image = preprocess_for_train(image,224,224)
        image_batch,label_batch = tf.train.batch([image,label],batch_size=batch_size,num_threads=64,capacity=capacity)
        label_batch = tf.reshape(label_batch,[batch_size])

        return image_batch,label_batch

    #独热编码
    def onehot(self,labels):
        n_sample = len(labels)
        n_class = max(labels) + 1
        onehot_labels = np.zeros((n_sample,n_class))
        onehot_labels[np.arange(n_sample),labels] = 1
        return onehot_labels

    
if __name__ == "__main__" :
    utils = U()
    #模型重新训练与保存
    startTime = time()
    batch_size = 32
    capacity = 256   #内存中存储的最大数据容量
    means = [123.68,116.779,103.939]  #RGB三通道所减均值

    #xs,ys = utils.get_file("./data/train/")   #获取图像列表和标签列表
    xs,ys = utils.get_file("D:/编程代码/python程序/迁移学习/dogs-vs-cats/train/")   #获取图像列表和标签列表
    image_batch,label_batch = utils.get_batch(xs,ys,224,224,batch_size,capacity)

    x = tf.placeholder(tf.float32,[None,224,224,3])
    y = tf.placeholder(tf.int32,[None,2])



    vgg = vgg16(x)
    fc8_finetuining = vgg.probs
    #损失函数
    loss_function = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=fc8_finetuining,labels=y))
    #选择优化器
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001).minimize(loss_function)

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    vgg.load_weights("D:/编程代码/python程序/迁移学习/vgg16_weights.npz",sess)
    saver = tf.train.Saver()
    #启动线程
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord,sess=sess)

    epoch_start_time = time()

    for i in range(1000):

        images,labels = sess.run([image_batch,label_batch])
        labels = utils.onehot(labels)

        sess.run(optimizer,feed_dict={x:images,y:labels})
        loss = sess.run(loss_function,feed_dict={x:images,y:labels})
        print("Now the loss is %f"% loss)

        epoch_end_time = time()
        print("Current epoch takes: ",(epoch_end_time - epoch_start_time))
        epoch_start_time = epoch_end_time

        if (i+1) % 500 == 0:
            saver.save(sess,os.path.join("D:/编程代码/python程序/迁移学习/model/","epoch {:06d}.ckpt".format(i)))
        print("----------Epoch %d is finished----------"%i)

    saver.save(sess,"D:/编程代码/python程序/迁移学习/model/")
    print("Optimization Finished!")


    duration = time() - startTime
    print("Train Finished takes:","{:.2f}".format(duration))
    coording.request_stop()   #通知其他线程关闭
    coord.join(threads)

    #预测
    means = [123.68,116.79,103.939]
    x = tf.placeholder(tf.float32,[None,224,224,3])

    sess = tf.Session()
    vgg = vgg16(x)
    fc8_finetuining = vgg.probs

    saver = tf.train.Saver()
    print("Model restoting……")
    saver.restore(sess,"D:/编程代码/python程序/迁移学习/model/")   #恢复最后保存的模型

    filepath = "D:/编程代码/python程序/迁移学习/dogs-vs-cats/test1/403.jpg"   #狗的图片
    #filepath = "D:/编程代码/python程序/迁移学习/dogs-vs-cats/test1/458.jpg"   #猫的图片

    img = imread(filepath,mode="RGB")
    img = imresize(img,(224,224))
    img = img.astype(np.float32)
    for c in range(3):
        img[:,:,c] -= means[c]
    prob = sess.run(fc8_finetuining,feed_dict={x:[img]})
    max_index = np.argmax(prob)


    if max_index == 0:
        print("This is a cat with possibility %.6f"%prob[:,0])
    else:
        print("This is a dog with possibility %.6f"%prob[:,1])

Model restoting……
This is a dog with possibility 0.993427

Model restoting……
This is a cat with possibility 0.998492