【迁移学习】猫狗数据分类案例(TensorFlow2)

1 导包

import tensorflow as tf # 版本2.3
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import glob # 文件图片路径读取
import os

2 数据处理

共2000张猫狗图片数据集，下载地址：Kaggle dc_2000

# 训练集
train_image_path = glob.glob(r'dc_2000/train/*/*.jpg') # 注意路径修改
np.random.shuffle(train_image_path) # 乱序
# 从train_image_path中 分离出'dog'或'cat'，制作0 1标签（cat为1，dog为0）
train_image_label = [int(p.split('/')[2] == 'cat') for p in train_image_path] # 根据路径实际情况修改：p.split('/')[2]

def load_preprosess_image(path, label):
    image = tf.io.read_file(path) # 根据路径读取图片
    image = tf.image.decode_jpeg(image, channels=3) # 图片编码
    image = tf.image.resize(image, [256, 256]) # resize
    image = tf.cast(image, tf.float32) # 类型转换
    image = image/255 # 归一化
    return image, label
# 创建训练集的Dataset
train_ds = tf.data.Dataset.from_tensor_slices((train_image_path, train_image_label))
AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.map(load_preprosess_image, num_parallel_calls=AUTOTUNE) # 对每一条数据应用load_preprosess_image方法
# 查看数据中的图片
for img, label in train_ds.take(1):
    plt.imshow(img)

BATCH_SIZE = 8 
train_count = len(train_image_path)

train_ds = train_ds.shuffle(300).batch(BATCH_SIZE)
train_ds = train_ds.prefetch(AUTOTUNE)

# 测试集 方法同上
test_image_path = glob.glob(r'dc_2000/test/*/*.jpg') # 注意路径修改
test_label = [int(p.split('/')[2] == 'cat') for p in test_image_path] # 同上，根据实际路径进行修改
test_ds = tf.data.Dataset.from_tensor_slices((test_image_path, test_label))
test_ds = test_ds.map(load_preprosess_image, num_parallel_calls=AUTOTUNE)
test_ds = test_ds.batch(BATCH_SIZE)

test_count = len(test_image_path) # 测试集数量

3 创建模型，使用VGG16作为预训练模型

covn_base = tf.keras.applications.VGG16(weights='imagenet',     # 使用了在imagenet中预训练好的权重
                                        include_top=False,      # 是否包含最后的全连接层
                                        input_shape=(256, 256, 3),
                                        pooling='avg')
# 如果上面VGG16下载失败，需要手动下载如下：
# vgg16_weights_tf_dim_ordering_tf_kernels.h5
# vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5
# 将其放入系统用户目录下的 .keras/models 文件夹中                                        
covn_base.trainable = False # 设置为不可训练
print(covn_base.summary())

输出：

Model: "vgg16"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 256, 256, 3)]     0         
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 256, 256, 64)      1792      
_________________________________________________________________
block1_conv2 (Conv2D)        (None, 256, 256, 64)      36928     
_________________________________________________________________
block1_pool (MaxPooling2D)   (None, 128, 128, 64)      0         
_________________________________________________________________
block2_conv1 (Conv2D)        (None, 128, 128, 128)     73856     
_________________________________________________________________
block2_conv2 (Conv2D)        (None, 128, 128, 128)     147584    
_________________________________________________________________
block2_pool (MaxPooling2D)   (None, 64, 64, 128)       0         
_________________________________________________________________
block3_conv1 (Conv2D)        (None, 64, 64, 256)       295168    
_________________________________________________________________
block3_conv2 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_conv3 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_pool (MaxPooling2D)   (None, 32, 32, 256)       0         
_________________________________________________________________
block4_conv1 (Conv2D)        (None, 32, 32, 512)       1180160   
_________________________________________________________________
block4_conv2 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_conv3 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_pool (MaxPooling2D)   (None, 16, 16, 512)       0         
_________________________________________________________________
block5_conv1 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv2 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv3 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_pool (MaxPooling2D)   (None, 8, 8, 512)         0         
_________________________________________________________________
global_average_pooling2d_1 ( (None, 512)               0         
=================================================================
Total params: 14,714,688
Trainable params: 0
Non-trainable params: 14,714,688
_________________________________________________________________

model = keras.Sequential() 
model.add(covn_base) # 添加VGG16
# 如下为自己的模型
model.add(layers.Dense(512, activation='relu'))
model.add(layers.Dense(1)) # 分类输出

4 训练模型

# Adam优化器，二元交叉熵损失 
model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.0005),
             loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['acc'])
# 训练15个批次，同时在测试集上验证准确率
history = model.fit(
                    train_ds,
                    steps_per_epoch=train_count//BATCH_SIZE,
                    epochs=15,
                    validation_data=test_ds,
                    validation_steps=test_count//BATCH_SIZE)
# 可视化，查看训练集和测试集上的准确率Acc
print(history.history.keys())
plt.plot(history.epoch, history.history.get('acc'), label='train_acc')
plt.plot(history.epoch, history.history.get('val_acc'), label='val_acc')
plt.legend()