在TensorFlow中使用自定义数据集训练自己的模型

最新推荐文章于 2024-09-17 13:19:51 发布

t01051

最新推荐文章于 2024-09-17 13:19:51 发布

阅读量6.2k

点赞数 10

文章标签： tensorflow 人工智能 python

本文链接：https://blog.csdn.net/t01051/article/details/121221619

版权

神经网络专栏收录该内容

1 篇文章 1 订阅

订阅专栏

在TensorFlow中使用自定义数据集训练自己的模型

写在前面的话

今年电赛终于结束了，身边不少小伙伴都选择了送药小车的题目，刚开始可能都觉得简单吧，循迹小车+数字识别就可以搞定。刚开始很多朋友都考虑使用OpenMv作为数字识别的平台，我手上除了大家津津乐道的OpenMv之外还有国产的K210。老实说就我个人而言我更愿意去使用K210。首先OpenMv作为一个开源项目已经很好的被K210支持，这意味着使用星瞳OpenMv能做的大部分操作K210都可以完成，甚至更好。其次，星瞳OpenMv主要基于ST公司的STM32系列处理器，目前最好的应该是用到了H7系列的片子。尽管它也能运行神经网络，但有一说一它的性能及其有限，一旦运行神经网络之后帧率降得挺严重的，某组小伙伴3 ~ 4帧的帧率。而国产芯片K210不仅价格较之更便宜，其算力更是高达1TOPS，同时硬件KPU支持神经网络常见层，性能强悍！比较有意思的是今年不少小伙伴都打算使用OpenMv作为数字识别的方案，结果纷纷翻车。倒也不是说OpenMv不能做，估计大家不是很熟悉没反应过来吧。星瞳OpenMv神经网络的在线训练平台虽然免费但是对训练时间还是有限制的，手动修改网络层数深一些之后会导致训练线程超时终止，不如本地TensorFlow训练来得实在。追根究底在OpenMv上跑神经网络基本都是tflite模型，TensorFlow本地训练好之后将模型转换好想来在OpenMv上跑也是没啥问题的。由于个人对K210的偏爱电赛期间我的OpenMv基本就放在旁边吃灰了，也看到过有几个小组和我一样使用K210做为数字识别的方案，但大家更多使用的是第三方提供的训练工具或者模型，训练后最终的效果大都不尽人意。影响卷积神经网络结果的原因有很多，模型的结构、优化器函数的选择、数据集的数量和质量等都会影响最终模型的效果。
本文仅介绍如何在TensorFlow上使用自己的数据集训练自己的模型，关于CNN模型的构建、调优以及转换为K210支持的kmodel格式等问题均不在本文讨论范围内，有兴趣的小伙伴可自行查阅相关文档~~

一、自定义数据集的目录结构（以今年电赛数字识别为例）

在这里插入图片描述
其中 Training 是训练集， Validation 是验证集

各类中的部分样本：
在这里插入图片描述

上面的数据集比较混乱，原因是刚开始我打算直接使用mnist数据集进行训练，结果识别效果并不是很好。重新制作数据集训练的时候直接使用的是前面训练的模型，虽然题目要求识别的数字只有 1~8 但 0 和 9 仍被保留了下来。

具体步骤：
    1.新建一个用于存放数据集的目录 dataset
    2.在 dataset 目录中创建 Training 目录用于存放训练集图片， Validation 目录则存放验证集图片（如果需要验证集的话）
    3.Training 目录与 Validation 目录中每个文件夹的名字就是一个标签 （Training目录与Validation目录中文件夹的数量及名称应保持一致） 。如上图中Training目录下的文件夹 0~9，分别对应mnist数据集中的 label 0~9
    4.dataset/Training/*目录下存放的就是训练集中对应该标签的图像数据（如上图中 各类中的部分样本 所示）

至此，数据集的目录已经构建完成。下一步则需要在TensorFlow中读入数据集的目录结构并解析、转换为tensor的数据类型以满足TensorFlow后面自动推理的要求

二、在TensorFlow中读取数据集

一般情况下，若使用TensorFlow自带的数据集可通过如下方式加载：

import tensorflow as tf

(train_data,train_label),(test_data,test_label) = tf.keras.datasets.mnist.load_data()

对于自定义数据集而言，我们有多种导入的方式，下文中将介绍如下两种：
2.1 纯手工打造
2.2 利用TensorFlow keras ImageDataGenerator

2.1 纯手工打造

先介绍第一种纯手工打造的方式。这种方式可以让你对整个数据集的读取、处理流程更清晰，灵活性更好。我们则需要手动去实现中间的过程，首先导入相关的包：

import tensorflow as tf
import pathlib

数据集的目录结构参照上文所述制作即可。自定义数据集的载入流程可以分成如下几个步骤：
    1.获取所有图片的路径
    2.获取标签并转换为数字
    3.读取图片并进行相应的预处理
    4.打包图片与标签

下面将以在TensorFlow中导入训练集 (dataset/Training) 为例进行说明，验证集 Validation 和测试集导入方式同理，下文将不再进行说明。

2.1.1 获取所有图片的路径

# 指定训练集数据的路径
my_dataset_path = 'dataset/Training'
# 指定图像要调整的大小，图像大小应与模型输入层保持一致
my_image_size = (32,32)
# 指定图像维度    1，单通道（例如灰度图）；3，三通道（彩图）
my_input_shape = my_image_size + (3,)
# 指定batch
my_batch = 32
# shuffle buffer size
my_shuffle_buffer_size = 1000

AUTOTUNE = tf.data.experimental.AUTOTUNE

# 获取所有图像文件的路径
dataset_path = pathlib.Path(my_dataset_path)
all_images_paths = [str(path) for path in list(dataset_path.glob('*/*'))]
print('所有文件的路径:', all_images_paths)
print('文件总数:', len(all_images_paths))

输出结果如下:

所有文件的路径: ['dataset\\Training\\0\\1.jpg', 'dataset\\Training\\0\\21.jpg', 'dataset\\Training\\1\\00000.jpg', 'dataset\\Training\\1\\00001.jpg', 'dataset\\Training\\2\\00000.jpg', 'dataset\\Training\\2\\00001.jpg', 'dataset\\Training\\3\\00027.jpg', 'dataset\\Training\\3\\00028.jpg', 'dataset\\Training\\4\\00000.jpg', 'dataset\\Training\\4\\00001.jpg', 'dataset\\Training\\5\\00161.jpg', 'dataset\\Training\\5\\00162.jpg', 'dataset\\Training\\6\\00001.jpg', 'dataset\\Training\\6\\00002.jpg', 'dataset\\Training\\7\\00163.jpg', 'dataset\\Training\\7\\00164.jpg', 'dataset\\Training\\8\\00000.jpg', 'dataset\\Training\\8\\00001.jpg', 'dataset\\Training\\9\\19.jpg', 'dataset\\Training\\9\\4.jpg']

此时，all_images_paths 列表中存储了训练数据集中所有图片的路径，在后面我们需要通过这些路径来将图片读入到内存中

2.1.2 获取标签并转换为数字

在目标检测的任务中通常会有多个标签 (label) ，为了方便人们区分，这些标签通常会使用字符串来进行描述。但是在TensorFlow进行推理的过程中是无法是用文本类型的标签的，我们需要将其转换成数字。
首先，我们需要通过解析图片路径 (all_images_paths) 的上层目录来获取标签：

# 获取标签名称
label_name = [i.name for i in dataset_path.iterdir() if i.is_dir()]
print('标签名称:', label_name)

输出结果如下:

标签名称: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

可以看到这些标签正好就是我们目录的名字，这证明我们程序的解析是正确的。接着，我们需要为这些标签分配一个唯一的数字，用这个数字来代替标签的名字：

# 因为训练时参数必须为数字，因此为标签分配数字索引
label_index = dict((name,index)for index,name in enumerate(label_name))
print('为标签分配数字索引:', label_index)

输出结果如下:

为标签分配数字索引: {'9': 9, '1': 1, '6': 6, '5': 5, '2': 2, '0': 0, '3': 3, '4': 4, '7': 7, '8': 8}

然后我们将图片与标签的数字索引进行配对，务必保证 all_images_paths 列表中图像数据的标签必须与数字索引一致：

# 将图片与标签的数字索引进行配对(number encodeing)
number_encodeing = [label_index[i.split('\\')[2]]for i in all_images_paths]
print('number_encodeing:', number_encodeing, type(number_encodeing))

输出结果如下:

number_encodeing: [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9] <class 'list'>

为了方便截图我删除了大量的样本，此处我们数据集dataset/Training中每一类都有两个图片样本，因此读取完成后的结果与上面的输出结果是一致的
在这里插入图片描述

为方便训练，多数情况下你还需要对标签做One Hot变换：

label_one_hot = tf.keras.utils.to_categorical(number_encodeing)
print(label_one_hot)

输出结果如下:

[[1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]]

2.1.3 读取图片并进行相应的预处理

在上面的步骤中我们完成了对数据集标签的处理，现在我们需要从磁盘中载入图像并在正式开始训练之前对图像进行相应的预处理：

def process(path, label):
    # 读入图片文件
    image = tf.io.read_file(path)
    # 将输入的图片解码为gray或者rgb
    image = tf.image.decode_jpeg(image, channels=my_input_shape[2])
    # 调整图片尺寸以满足网络输入层的要求
    image = tf.image.resize(image, my_image_size)
    # 归一化
    image /= 255.
    return image,label

process方法接收一个路径 （还有一个标签，为了方便后面打包），在该方法的内部使用TensorFlow的一些方法完成对图像的读取，同时调整图像的尺寸和维度以便满足模型的需要。最后对调整好的图像数据进行归一化处理，一般情况下我们将范围映射到 0~1 的区间，当然你也可以映射到 -0.5~0.5 等区间，根据自己的需要去设计即可。一个好的区间可以让模型更好的收敛同时精度也得以提升。
使用matplotlib.pyplot查看图像：

import matplotlib.pyplot as plt

img = process('dataset\\Training\\1\\00001.jpg', 1)
# 或
img = process(all_images_paths[2], 1)

print(img[0].shape)
plt.imshow(img[0])
plt.show()

输出结果：

(32, 32, 3)

在这里插入图片描述

从输出的结果中可以看到图像的尺寸已经调整到与我们设置的 my_imput_shape=(32,32,3) 一致

2.1.4 打包图片与标签

# 将数据与标签拼接到一起
label_one_hot = tf.cast(label_one_hot, tf.int32)
path_ds = tf.data.Dataset.from_tensor_slices((all_images_paths, label_one_hot))
image_label_ds = path_ds.map(process, num_parallel_calls=AUTOTUNE)
print('image_label_ds:', image_label_ds)

输出结果：

image_label_ds: <ZipDataset shapes: ((32, 32, 3), (10,)), types: (tf.float32, tf.int32)>

tf.data.Dataset.from_tensor_slices 的作用是对数据集进行切片。上文中我们将所有图片的路径 all_images_paths 以及 label_one_hot 传入，经 map() 后得到一个二元组，其中索引0存放的是图像数据，索引1存放的则是其label对应的One Hot编码:

# 以第3个数据为例，输出其结果
res = [i for i in image_label_ds.take(3)][-1]
print(res[0])
print(res[1])

# 显示图片
plt.imshow(res[0])
plt.show()

输出结果：

tf.Tensor(
[[[0.48088235 0.54362744 0.49803922]
  [0.5034314  0.5495098  0.5122549 ]
  [0.5264706  0.56960785 0.5382353 ]
  ...
  [0.54068625 0.60343134 0.5602941 ]
  [0.5112745  0.59656864 0.5254902 ]
  [0.4970588  0.58137256 0.53137255]]

 [[0.46813726 0.53088236 0.49019608]
  [0.50980395 0.5529412  0.52156866]
  [0.5137255  0.5568628  0.5254902 ]
  ...
  [0.5352941  0.5980392  0.55490196]
  [0.5132353  0.57598037 0.5328431 ]
  [0.5156863  0.5745098  0.5470588 ]]

 [[0.49509802 0.5421569  0.5029412 ]
  [0.5107843  0.55784315 0.51862746]
  [0.5323529  0.5715686  0.54019606]
  ...
  [0.5387255  0.5857843  0.54656863]
  [0.50980395 0.5715686  0.5205882 ]
  [0.50735295 0.57009804 0.53088236]]

 ...

 [[0.38333333 0.44607842 0.40686274]
  [0.42009804 0.4632353  0.43186274]
  [0.43333334 0.4764706  0.46078432]
  ...
  [0.422549   0.50490195 0.44607842]
  [0.4240196  0.5181373  0.44558823]
  [0.41813725 0.5122549  0.45735294]]

 [[0.3882353  0.4392157  0.4       ]
  [0.40882352 0.4598039  0.42058823]
  [0.42941177 0.48039216 0.44509804]
  ...
  [0.44166666 0.5122549  0.45735294]
  [0.42647058 0.5088235  0.44215685]
  [0.39950982 0.4779412  0.43088236]]

 [[0.36862746 0.44509804 0.39803922]
  [0.3882353  0.46470588 0.41764706]
  [0.40441176 0.47990197 0.4377451 ]
  ...
  [0.43186274 0.502451   0.44754902]
  [0.41127452 0.49362746 0.4269608 ]
  [0.40686274 0.4852941  0.43823528]]], shape=(32, 32, 3), dtype=float32)
tf.Tensor([0 1 0 0 0 0 0 0 0 0], shape=(10,), dtype=int32)

在这里插入图片描述

为防止过拟合增强模型的泛化能力，我们还需要将数据集中的顺序打乱，先来看看打乱之前图片的顺序：

def display_more_image(image_label_ds, s_pos, e_pos, max_r, max_c):
    if e_pos >= s_pos:
        index = 1
        plt.figure()
        for n,image in enumerate(image_label_ds.take(e_pos)):
            if n >= s_pos-1:
                img = image[0]
                label = image[1]
                plt.subplot(max_r, max_c, index)
                index += 1
                plt.imshow(img)
                plt.xlabel(str(list(label.numpy()).index(max(label.numpy()))))
        plt.show()

display_more_image(image_label_ds, 1, 6, 2, 3)

输出结果：
在这里插入图片描述

现在打乱顺序

# 打乱dataset中的元素并设置batch
image_label_ds = image_label_ds.shuffle(my_shuffle_buffer_size).batch(my_batch)

显示打乱顺序后的图片：

# 显示第一个bath的前六张图片
index = 1
plt.figure()
for i in image_label_ds.take(1):
    for j in range(6):
        plt.subplot(2, 3, index)
        index += 1
        # print(i[1][j])
        plt.imshow(i[0][j])
        plt.xlabel(str(list(i[1][j].numpy()).index(max(i[1][j].numpy()))))

# 务必等上面的for执行完成后再调用显示
plt.show()

执行结果：
在这里插入图片描述

可见使用shuffle打乱顺序后的输出与打乱之前的图片顺序是不一致的！

至此，TensorFlow加载自定义数据集的操作就完成了，之后就是构建网络模型以及训练。

2.2 利用TensorFlow keras ImageDataGenerator

这种导入数据的方式较之上面那种更简单、更优雅，主要使用keras ImageDataGenerator，你只需要几行代码即可完成数据集的导入：

my_image_size = (32,32)
my_input_shape = my_image_size + (3,)
# 指定训练次数
my_train_epochs = 2
# 指定batch
my_batch = 32
# 创建
两个数据生成器，指定scaling范围0~1
train_datagen = ImageDataGenerator(rescale=1/255)
# validation_datagen = ImageDataGenerator(rescale=1/255)

# 指向训练数据文件夹
train_generator = train_datagen.flow_from_directory(
    './dataset/Training',           # 训练数据所在文件夹
    target_size=my_image_size,         # 指定输出尺寸
    batch_size=my_batch,
    color_mode=  'rgb' if my_input_shape[2]==3 else 'grayscale',
    class_mode='categorical')            # 指定分类# 创建两个数据生成器，指定scaling范围0~1
train_datagen = ImageDataGenerator(rescale=1/255)
# validation_datagen = ImageDataGenerator(rescale=1/255)

# 指向验证数据文件夹
# validation_generator = validation_datagen.flow_from_directory(
#     './dataset/Validation',
#     target_size=my_image_size,
#     batch_size=my_batch,
#     color_mode=  'rgb' if my_input_shape[2]==3 else 'grayscale',
#     class_mode='categorical')

三、验证自定义数据集是否可用

下面以我今年电赛中数字识别的网络模型为例测试上文中自定义的数据集是否可用，关于模型的构建不是本文重点，有兴趣的小伙伴可自行查阅相关资料，下文直接给出完整的测试用例：

3.1 验证纯手工方式加载的数据

#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Python version: 3.5.4
'''
@File    :   tf_data_generator.py
@Time    :   2021/11/04 16:11:19
@Author  :   Wu Xueru 
@Version :   1.0
@Contact :   t01051@163.com
@License :   
@Desc    :   
'''

# here put the import lib
import tensorflow as tf
import pathlib
import matplotlib.pyplot as plt
from time import strftime


my_dataset_path = 'dataset/Training'
my_image_size = (32,32)
my_input_shape = my_image_size + (3,)
# 指定训练次数
my_train_epochs = 2
# 指定batch
my_batch = 32
# shuffle buffer size
my_shuffle_buffer_size = 1000

AUTOTUNE = tf.data.experimental.AUTOTUNE

# 获取所有文件路径
dataset_path = pathlib.Path(my_dataset_path)
all_images_paths = [str(path) for path in list(dataset_path.glob('*/*'))]
print('所有文件的路径:', all_images_paths)
print('文件总数:', len(all_images_paths))

# 获取标签名称
label_name = [i.name for i in dataset_path.iterdir() if i.is_dir()]
print('标签名称:', label_name)
# 因为训练时参数必须为数字，因此为标签分配数字索引
label_index = dict((name,index)for index,name in enumerate(label_name))
print('为标签分配数字索引:', label_index)

# 将图片与标签的数字索引进行配对(number encodeing)
number_encodeing = [label_index[i.split('\\')[2]]for i in all_images_paths]
print('number_encodeing:', number_encodeing, type(number_encodeing))
label_one_hot = tf.keras.utils.to_categorical(number_encodeing, num_classes=10)
print('label_one_hot:', label_one_hot)


def process(path,label):
    # 读入图片文件
    image = tf.io.read_file(path)
    # 将输入的图片解码为gray或者rgb
    image = tf.image.decode_jpeg(image, channels=my_input_shape[2])
    # 调整图片尺寸以满足网络输入层的要求
    image = tf.image.resize(image, my_image_size)
    # 归一化
    image /= 255.
    return image,label

# 将数据与标签拼接到一起
path_ds = tf.data.Dataset.from_tensor_slices((all_images_paths, tf.cast(label_one_hot, tf.int32)))
image_label_ds = path_ds.map(process, num_parallel_calls=AUTOTUNE)
print('image_label_ds:', image_label_ds)
steps_per_epoch=tf.math.ceil(len(all_images_paths)/my_batch).numpy()
print('steps_per_epoch', steps_per_epoch)

# 打乱dataset中的元素并设置batch
image_label_ds = image_label_ds.shuffle(my_shuffle_buffer_size).batch(my_batch)


if __name__ == '__main__':
    # 定义模型
    # 输入层
    input_data = tf.keras.layers.Input(shape=my_input_shape)
    # 第一层
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(input_data)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)
    # 第二层
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)
    # 第三层
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)

    # 铺平
    dense = tf.keras.layers.Flatten()(middle)
    dense = tf.keras.layers.Dropout(0.1)(dense)
    dense = tf.keras.layers.Dense(60, activation='relu')(dense)
    # 输出
    # 输出层
    output_data = tf.keras.layers.Dense(len(label_name), activation='softmax')(dense)
    # 确认输入位置和输出位置
    model = tf.keras.Model(inputs=input_data, outputs=output_data)

    # 定义模型的梯度下降和损失函数
    model.compile(optimizer=tf.optimizers.Adam(1e-4), 
                loss=tf.losses.categorical_crossentropy,
                metrics=['accuracy'])

    # 打印模型结构
    model.summary()

    # 开始训练
    start_time = strftime("%Y-%m-%d %H:%M:%S")
    history = model.fit(
        image_label_ds,
        epochs=my_train_epochs,
        verbose=1,
        steps_per_epoch=int(steps_per_epoch))

    end_time = strftime("%Y-%m-%d %H:%M:%S")
    print('开始训练的时间：', start_time)
    print('结束训练的时间：', end_time)

输出结果：
在这里插入图片描述

从输出中可以看到自定义的数据集已经被正确读取，TensorFlow能够正常进行推理。

3.1 验证利用TensorFlow keras ImageDataGenerator方式加载的数据

# here put the import lib
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
from time import strftime
from os import path

my_image_size = (32,32)
my_input_shape = my_image_size + (3,)
# 指定训练次数
my_train_epochs = 2
# 指定batch
my_batch = 32
# shuffle buffer size
my_shuffle_buffer_size = 1000
save_model_path = path.dirname(path.abspath(__file__))

# 创建两个数据生成器，指定scaling范围0~1
train_datagen = ImageDataGenerator(rescale=1/255)
# validation_datagen = ImageDataGenerator(rescale=1/255)

# 指向训练数据文件夹
train_generator = train_datagen.flow_from_directory(
    './dataset/Training',           # 训练数据所在文件夹
    target_size=my_image_size,         # 指定输出尺寸
    batch_size=my_batch,
    color_mode=  'rgb' if my_input_shape[2]==3 else 'grayscale',
    class_mode='categorical')            # 指定分类

# 指向验证数据文件夹
# validation_generator = validation_datagen.flow_from_directory(
#     './dataset/Validation',
#     target_size=my_image_size,
#     batch_size=my_batch,
#     color_mode=  'rgb' if my_input_shape[2]==3 else 'grayscale',
#     class_mode='categorical')

if __name__ == '__main__':
    # 定义模型
    # 输入层
    input_data = tf.keras.layers.Input(shape=my_input_shape)
    # 第一次卷积
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(input_data)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)
    # 第二次卷积
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)
    # 第三次卷积
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.Conv2D(128, kernel_size=[3,3], strides=(1,1), padding='same', activation=tf.nn.relu)(middle)
    middle = tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2, padding='same')(middle)

    # 铺平
    dense = tf.keras.layers.Flatten()(middle)
    dense = tf.keras.layers.Dropout(0.1)(dense)
    dense = tf.keras.layers.Dense(60, activation='relu')(dense)
    # 输出
    # 输出层
    output_data = tf.keras.layers.Dense(10, activation='softmax')(dense)
    # 确认输入位置和输出位置
    model = tf.keras.Model(inputs=input_data, outputs=output_data)

    # 定义模型的梯度下降和损失函数
    model.compile(optimizer=tf.optimizers.Adam(1e-4), 
                loss=tf.losses.categorical_crossentropy,
                metrics=['accuracy'])

    # 打印模型结构
    model.summary()

    # 开始训练
    start_time = strftime("%Y-%m-%d %H:%M:%S")
    history = model.fit(
        train_generator,
        epochs=my_train_epochs,
        verbose=1)

    end_time = strftime("%Y-%m-%d %H:%M:%S")
    print('开始训练的时间：', start_time)
    print('结束训练的时间：', end_time)