Tensorflow2+版本-猫狗分类完整训练预测流程

一、下载数据集

数据集我放在百度网盘中：百度网盘 请输入提取码

提取码：itai

为了不修改后文代码就可以直接运行，请把数据集解压路径改成下图所示（即ipynb文件需要同train和test文件在同一目录下）：

二、导入所需库（没有的直接pip）

import glob
import warnings
import os, shutil
import numpy as np 
import tensorflow as tf
import keras.backend as K
from PIL import Image
from tensorflow import keras
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")
from keras.preprocessing import image
from keras.callbacks import ModelCheckpoint
from tensorflow.keras.callbacks import TensorBoard
from sklearn.model_selection import train_test_split
from tensorflow.keras.preprocessing.image import ImageDataGenerator

三、划分训练测试集

        划分完成后可以看到在all_data下分别有2个文件一个为训练数据集一个为测试数据集，其中训练集包含了10000张分好类别的猫狗图片，测试集下有2500张分好类别的猫狗图片。

# 下载的kaggle数据集路径
original_dataset_dir = './train' 
# 训练测试数据集放置路径
base_dir = './all_data' 

os.makedirs(base_dir)
train_dir = os.path.join(base_dir, 'train')
os.mkdir(train_dir)
validation_dir = os.path.join(base_dir, 'validation')
os.mkdir(validation_dir)

train_cats_dir = os.path.join(train_dir, 'cats')
os.mkdir(train_cats_dir)
train_dogs_dir = os.path.join(train_dir, 'dogs')
os.mkdir(train_dogs_dir)
validation_cats_dir = os.path.join(validation_dir, 'cats')
os.mkdir(validation_cats_dir)
validation_dogs_dir = os.path.join(validation_dir, 'dogs')
os.mkdir(validation_dogs_dir)

fnames = ['cat.{}.jpg'.format(i) for i in range(10000)]
for fname in fnames:
    src = os.path.join(original_dataset_dir, fname)
    dst = os.path.join(train_cats_dir, fname)
    shutil.copyfile(src, dst)

fnames = ['cat.{}.jpg'.format(i) for i in range(10000, 12500)]
for fname in fnames:
    src = os.path.join(original_dataset_dir, fname)
    dst = os.path.join(validation_cats_dir, fname)
    shutil.copyfile(src, dst)

fnames = ['dog.{}.jpg'.format(i) for i in range(10000)]
for fname in fnames:
    src = os.path.join(original_dataset_dir, fname)
    dst = os.path.join(train_dogs_dir, fname)
    shutil.copyfile(src, dst)

fnames = ['dog.{}.jpg'.format(i) for i in range(10000, 12500)]
for fname in fnames:
    src = os.path.join(original_dataset_dir, fname)
    dst = os.path.join(validation_dogs_dir, fname)
    shutil.copyfile(src, dst)

print('total training cat images:', len(os.listdir(train_cats_dir)))
print('total training dog images:', len(os.listdir(train_dogs_dir)))
print('total test cat images:', len(os.listdir(validation_dogs_dir)))
print('total test dog images:', len(os.listdir(validation_dogs_dir)))

四、搭建模型。

        构建模型并查看模型架构和参数，此处采用卷积和最大池化操作，后续在全连接层加入一层dropout防止过拟合。

model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(64, 64, 3)),
    tf.keras.layers.MaxPooling2D(2,2),
    
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    
    tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    
    #转换为一维
    tf.keras.layers.Flatten(),
    
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(1, activation='sigmoid'),
])

model.summary()

五、定义优化器和损失函数。

        由于是二分类问题所以loss选择binary_crossentropy。

#定义优化器和损失函数
model.compile(
    optimizer=tf.keras.optimizers.Adam(),
    loss='binary_crossentropy',
    metrics=['acc']
)

六、数据增强。

        为防止过拟合，对训练集的图片数据进行数据增强操作，此处用到了归一化、随机翻转、平移、随机裁剪和随机缩放，对测试集的图片只进行归一化操作即可。

#训练集
train_datagen = ImageDataGenerator(
    rescale=1/255,
    #随机翻转
    rotation_range=40,
    #平移
    width_shift_range=0.2,
    height_shift_range=0.2,
    #随机裁剪
    shear_range=0.2,
    #随机缩放
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest'
)

#测试集
validation_datagen = ImageDataGenerator(
    rescale=1/255
)

七、构建训练和验证数据的生成器。

        不断从文件夹中读取图片数据减少内存的消耗。

train_generator = train_datagen.flow_from_directory(
        train_dir,
        #修改图片大小为64*64
        target_size=(64,64),
        #一次训练20张图片
        batch_size=20,
        #二分类使用binary 
        class_mode='binary'
)

validation_generator = validation_datagen.flow_from_directory(
        validation_dir,
        target_size=(64,64),
        batch_size=20,
        class_mode='binary'
)

八、训练。

        构造tensorboard可视化画板以及checkpoint保存最佳模型，开始训练。

#存放logs
summary_writer = TensorBoard(log_dir="./logs", histogram_freq=1)

#保存最佳模型
base_dir = './model'
if not os.path.exists(base_dir):
    os.makedirs(base_dir)
filepath='./model/weights.best.hdf5'
# 有一次提升, 则覆盖一次.
checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1,save_best_only=True,mode='max') 


# fit_generator相当于一个生成器，动态产生所需batch数据
history = model.fit(
    train_generator,
    #给定停止条件  10000/20   image/batch_size
    steps_per_epoch=500,
    epochs=60,
    validation_data=validation_generator,
    #2500/20
    validation_steps=125,
    verbose=2,
    callbacks=[summary_writer, checkpoint],
)

        训练时使用tensorboard查看loss和acc变化，在pycharm终端输入以下代码，然后点击端口进入浏览器即可：

tensorboard --logdir=./logs

九、预测。

预测单张图片

image = Image.open('./test/300.jpg')
plt.imshow(image)
img = np.array(tf.image.resize(np.array(image), (64,64)))
img = np.reshape(img, (1,64,64,3))
img = img/255
predict = model.predict(img)
print(predict)
if predict>0.5:
    print('This is dog.')
else:
    print('This is cat.')

注：如果想预测多张图片，可以添加test的生成器来预测。