欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

TensofFlow制作自己的数据集,并训练CNN网络

程序员文章站 2024-03-19 18:31:58
...

前面的几篇文章我们学习了使用Tensorflow搭建简单拟合神经网络、CNN、RNN和自编码等。但是在那些简单的例子里面我们都是使用的标准的MNIST数据集,而大部分时间我们更需要使用自己的数据集训练神经网络。因此,在这里我们介绍如何使用TFrecords处理自己的数据集并训练CNN网络。我们以kaggle比赛中的猫狗大战数据集为例。

(1)下载并分类处理猫狗大战训练集

猫狗大战数据集下载链接如下所示

 https://pan.baidu.com/s/13hw4LK8ihR6-6-8mpjLKDA

密码:dmp4

下载后得到数据的训练集和测试集

TensofFlow制作自己的数据集,并训练CNN网络

在这里我们只需要训练集,训练集中包括猫狗图像25000张,数据太多了而且猫狗的图像都在一个文件夹中,只于图像的名字区分,很不利于我们制作训练集啊,因此我们写了一个简单的程序从训练集中取1000张图像,并将猫狗分别保存到文件夹内,对了为了方便大家单独新建个文件夹用于保存分类后的图像,可以参考下面代码中的路径。

import os
import shutil

#下载得到的训练集图像
image_path='/home/cyy/python_code/tensorflow/kaggle/train'
#将猫狗分类保存的路径
train_path='/home/cyy/python_code/tensorflow/dog_cat_classification/data'

image_list=os.listdir(image_path)
#读取1000张猫狗图像,按照图像名字分别保存
for image_name in image_list[0:1000]:
    class_name=image_name[0:3]
    save_path=os.path.join(train_path,class_name)
    if not(os.path.exists(save_path)):
        os.mkdir(save_path)
    file_name=os.path.join(image_path,image_name)
    save_name=os.path.join(save_path,image_name)
    shutil.copyfile(file_name,save_name)

这样就将1000张猫狗图像分别保存到了两个文件夹中,如下图所示

TensofFlow制作自己的数据集,并训练CNN网络

(2)制作tfrecords文件

代码起名为make_own_data.py
tfrecord会根据你选择输入文件的类,自动给每一类打上同样的标签。

代码如下:

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

#存放数据文件夹
data_path='/home/cyy/python_code/tensorflow/dog_cat_classification/data'
classes={'dog','cat'}#两个类名
#保存的文件名
data_file=tf.python_io.TFRecordWriter('dog_and_cat_train.ftrecords')
#分别读取两个类并生成索引
for index,class_name in enumerate(classes):
    class_path=os.path.join(data_path,class_name)
    #分别读取每个类下的图像
    for im_name in os.listdir(class_path):
        image_path=os.path.join(class_path,im_name)
        
        img=Image.open(image_path)
        img=img.resize((128,128))#将图像resize成128*128
        img_raw=img.tobytes()#生成二进制
        example=tf.train.Example(features=tf.train.Features(feature={
            'label':tf.train.Feature(int64_list=tf.train.Int64List(value=[index])),
            'img_raw':tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
            }))#example对象对label和image数据进行封装
        data_file.write(example.SerializeToString())#序列化为字符串
data_file.close()

运行完后在代码目录下会生成文件如下所示:

TensofFlow制作自己的数据集,并训练CNN网络

(3)读取数据程序

将图片和标签读出,图片reshape为128x128x3。
读取代码单独作为一个文件,起名为ReadMyOwnData.py

代码如下:

import tensorflow as tf

def read_and_decode(filename):#读入tfrecords文件
    filename_queue=tf.train.string_input_producer([filename])
    reader=tf.TFRecordReader()
    _,serialized_example=reader.read(filename_queue)#生成一个queue队列
    #将image数据和label数据读取出来放入features
    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,[128,128,3])#将图像reshape成128*128
    img=tf.cast(img,tf.float32)*(1./255)-0.5#将image归一化
    label=tf.cast(features['label'],tf.int32)#将label转化成int32
    return(img,label)

 

 (4)搭建CNN网络并完成训练

要把我们读取文件的ReadMyOwnData导入,网络进行两次卷积操作,两次最大池化,**函数ReLU,全连接层,最后y_conv是softmax输出的二类问题。损失函数用交叉熵,优化算法Adam。

代码如下:

import tensorflow as tf
import numpy as np
import ReadMyOwnData

batch_size=50

#定义初始化权重和偏置函数
def weight_variable(shape):
    return(tf.Variable(tf.random_normal(shape,stddev=0.01)))
def bias_variable(shape):
    return(tf.Variable(tf.constant(0.1,shape=shape)))
#定义输入数据和dropout占位符
X=tf.placeholder(tf.float32,[batch_size,128,128,3])
y_=tf.placeholder(tf.float32,[batch_size,1])
keep_pro=tf.placeholder(tf.float32)

#搭建网络
def model(X,keep_pro):
    w1=weight_variable([5,5,3,32])
    b1=bias_variable([32])
    conv1=tf.nn.relu(tf.nn.conv2d(X,w1,strides=[1,1,1,1],padding='SAME')+b1)
    pool1=tf.nn.max_pool(conv1,ksize=[1,4,4,1],strides=[1,4,4,1],padding='SAME')
    
    w2=weight_variable([5,5,32,64])
    b2=bias_variable([64])
    conv2=tf.nn.relu(tf.nn.conv2d(pool1,w2,strides=[1,1,1,1],padding='SAME')+b2)
    pool2=tf.nn.max_pool(conv2,ksize=[1,4,4,1],strides=[1,4,4,1],padding='SAME') 
    tensor=tf.reshape(pool2,[batch_size,-1])
    dim=tensor.get_shape()[1].value
    w3=weight_variable([dim,1024])
    b3=bias_variable([1024])
    fc1=tf.nn.relu(tf.matmul(tensor,w3)+b3)
    h_fc1=tf.nn.dropout(fc1,keep_pro)
    w4=weight_variable([1024,2])
    b4=bias_variable([2])
    y_conv=tf.nn.softmax(tf.matmul(h_fc1,w4)+b4)
    return(y_conv)

#定义网络,并设置损失函数和训练器
y_conv=model(X,keep_pro)
cost=tf.reduce_mean(-tf.reduce_sum(y_*tf.log(y_conv),reduction_indices=[1]))
train_step=tf.train.AdamOptimizer(0.001).minimize(cost)
#计算准确率
correct_prediction=tf.equal(tf.argmax(y_conv,1),tf.argmax(y_,1))
accuracy=tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

#读取tfrecords数据
image,label=ReadMyOwnData.read_and_decode("dog_and_cat_train.ftrecords")
#定义会话,并开始训练
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    #定义多线程
    coord=tf.train.Coordinator()
    threads=tf.train.start_queue_runners(coord=coord)
    #定义训练图像和标签
    example=np.zeros((batch_size,128,128,3))
    l=np.zeros((batch_size,1))
    try:
        for i in range(20):
            #将数据存入example和l
            for epoch in range(batch_size):
                example[epoch],l[epoch]=sess.run([image,label])
            #开始训练
            sess.run(train_step,feed_dict={X:example,y_:l,keep_pro:0.5})
            print('train step','%04d ' %(i+1),'Accuracy=',sess.run(accuracy,feed_dict={X:example,y_:l,keep_pro:0.5}))
    except tf.errors.OutOfRangeError:
        print('done!')
    finally:
        coord.request_stop()
    coord.join(threads)

输出结果如下所示(代码在Python3下运行,如果报错尝试删除中文注释):

TensofFlow制作自己的数据集,并训练CNN网络

我们训练了20次,这是20个batch输出的结果,我们可以看到因为网络搭建的比较简单所以结果并不是特别好,如果想要更好的结果需要搭建更复杂的网络,当然计算量也会加大。