TensorFlow之VGG16实现猫狗大战

一、VGG16权重下载

链接：https://pan.baidu.com/s/1MumNcPZs5oyuUv28tls9eA
提取码：i3re

二、VGG16网络实现及加载权重/model.py

import tensorflow as tf
import numpy as np

class VGG16:
    def __init__(self,imgs):
        self.parameters=[] #载入后的权重
        self.imgs=imgs
        self.convlayers()
        self.fc_layers()
        self.probs=self.fc8

    def saver(self):
        return tf.train.Saver()

    def conv(self,name, input_data, out_channel):
        in_channel = input_data.get_shape()[-1]
        with tf.variable_scope(name):
            kernel = tf.get_variable("weights", [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 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 convlayers(self):
        #conv1
        self.conv1_1=self.conv("conv1re_1",self.imgs,64)
        self.conv1_2=self.conv("conv1_2",self.conv1_1,64)
        self.pool1=self.maxpool("poolre1",self.conv1_2)

        #cov2
        self.conv2_1=self.conv("conv2_1",self.pool1,128)
        self.conv2_2=self.conv("conv2_2",self.conv2_1,128)
        self.pool2=self.maxpool("pool2",self.conv2_2)

        #conv3
        self.conv3_1=self.conv("conv3_1",self.pool2,256)
        self.conv3_2=self.conv("conv3_2",self.conv3_1,256)
        self.conv3_3=self.conv("conv3-3",self.conv3_2,256)
        self.pool3=self.maxpool("pool3",self.conv3_3)

        #conv4
        self.conv4_1=self.conv("conv4_1",self.pool3,512)
        self.conv4_2=self.conv("conv4_2",self.conv4_1,512)
        self.conv4_3=self.conv("conv4_3",self.conv4_2,512)
        self.pool4=self.maxpool("pool4",self.conv4_3)

        #conv5
        self.conv5_1=self.conv("conv5_1",self.pool4,512)
        self.conv5_2=self.conv("conv5_2",self.conv5_1,512)
        self.conv5_3=self.conv("conv5_3",self.conv5_2,512)
        self.pool5=self.maxpool("pool5",self.conv5_3)

    def fc_layers(self):
        self.fc6=self.fc("fc6",self.pool5,4096,trainable=False)
        self.fc7=self.fc("fc7",self.fc6,4096,trainable=False)
        self.fc8=self.fc("fc8",self.fc7,2)
	
	#权重加载
    def load_weights(self,weight_file,sess):
        weights=np.load(weight_file)
        keys=sorted(weights.keys())
        print(keys)
        print(len(keys))
        for i,k in enumerate(keys):
            if i not in [30,31]: #剔除全连接层的权重，加载到parameters中
                sess.run(self.parameters[i].assign(weights[k]))
        print("all done")

三、数据集加载/creat_and_read.py

import tensorflow as tf
import  numpy as np
import os

def get_file(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])

    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


def get_batch(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=tf.image.resize_image_with_crop_or_pad(image,img_width,img_height)
    image=tf.image.per_image_standardization(image)
    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 one_hot(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

四、进行训练/train.py

import numpy as np
import tensorflow as tf
import model as model
import creat_and_read as reader

if __name__=='__main__':
    X_train,Y_train=reader.get_file('dogs_vs_cats')
    image_batch,label_batch=reader.get_batch(X_train,Y_train,224,224,25,256)
    x_imgs=tf.placeholder(tf.float32,[None,224,224,3])
    y_imgs=tf.placeholder(tf.int32,[None,2])
    vgg=model.VGG16(x_imgs)
    fc3_cat_and_dog=vgg.probs
    loss=tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=fc3_cat_and_dog,labels=y_imgs))
    optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.001).minimize(loss)
    sess=tf.Session()
    sess.run(tf.global_variables_initializer())
    vgg.load_weights('vgg16\\vgg16_weights.npz',sess)
    saver=vgg.saver()
    coord=tf.train.Coordinator()
    threads=tf.train.start_queue_runners(coord=coord,sess=sess)
    for i in range(20):
        image,label=sess.run([image_batch,label_batch])
        labels=reader.one_hot(label)
        sess.run(optimizer,feed_dict={x_imgs:image,y_imgs:labels})
        loss_record=sess.run(loss,feed_dict={x_imgs:image,y_imgs:labels})
        print("now loss is %f"%loss_record)

    saver.save(sess,'.\\vgg_finetuning_model\\')
    coord.request_stop()
    coord.join(threads)

五、测试集检测/test.py

import tensorflow as tf
import cv2
import model as model

imgs=tf.placeholder(tf.float32,[None,224,224,3])
sess=tf.Session()
vgg=model.VGG16(imgs)
fc3_cat_and_dog=vgg.probs
saver=vgg.saver()
saver.restore(sess,".\\vgg_finetuning_model\\")

import os
for root,sub_folders,files in os.walk("test1"):
    i=0
    cat=0
    dog=0
    for name in files:
        i+=1
        filepath=os.path.join(root,name)
        img1=[]
        try:
            img1=cv2.imread(filepath)
            img1=cv2.resize(img1,(224,224))
        except:
            print("remove",filepath)
        prob=sess.run(fc3_cat_and_dog,feed_dict={vgg.imgs:[img1]})[0]
        import numpy as np
        max_index=np.argmax(prob)
        if max_index==0:
            cat+=1
        else:
            dog+=1
        if i%50==0:
            acc=(cat*1.)/(cat+dog)
            print(acc)
            print("-----img num is %d-----"%i)