使用TensorFlow快速实现图片分类
前言
现在网络上深度学习非常热门,同时深度学习的门槛也越来越低。在TensorFlow中也提供了一些非常简单的方式让初学者甚至是没有接触过深度学习的小白都能快速上手
本文通过一个简单的图片分类例子来演示一下如何快速使用TensorFlow搭建完成目标
由于一些小型数据集(例如MNIST)不能表现出实际训练中遇到的一些问题,本文使用Kaggle上的nsfw-images-data 数据集来制作一个图片分类Bot,(实际使用案例 https://mirai.mamoe.net/topic/694 )
安装
自己去看另一个文章
数据的加载和预处理
在TensorFlow的官方快速上手教程里,数据集将直接完全载入到内存中。但对于本文的 30GB 大小的数据集来说这完全不限时。只能依靠tf.data
提供的数据加载流水线完成数据的预加载和处理。
使用TensorFlow提供的数据加载处理流水线还有一个好处:它能够绕过Python GIL的限制来加速数据读入和处理,这可以保证在任何时候io和预处理都不会成为影响训练时间的瓶颈
处理标签
数据集已经做好了测试集/训练集
的划分,图片总共有五个分类。我们需要先将图片标签使用One-Hot
的方式进行编码
1 | One-Hot编码就像这样 |
得到所有的分类并进行One-Hot编码后放入一个字典备用
1 | label = dict((v.name, k) for k, v in enumerate(pathlib.Path("setu/train").iterdir())) |
tf.one_hot
可以对整个列表里面的数据都进行One-Hot编码,只需要指定depth(分类数量)即可
得到整个数据集内的图片文件列表并打乱顺序
1 | train_files = [i for i in pathlib.Path("setu/train").glob("*/*")] |
现在开始构造tf.data.DataSet
对象,这个对象用于在训练或者是验证过程中向模型输送数据,同时可以添加数据预处理步骤
1 |
|
处理图片
现在开始处理图片了,我们可以编写一个图片变换函数来对训练数据集进行随机变换,以增加输入到模型中数据的多样性
1 | def _process_train(path): |
需要注意的是,稍后调用这个函数的输入为Tensor
格式,所以你可能无法在这个函数内使用一些基于numpy
或者是OpenCV
的数据预处理方案。同时使用tf提供的预处理器也能够保证性能
对测试数据集进行相同操作,在测试数据集中不要进行随机变换
1 | def _process_test(path): |
构建数据处理流水线
在上面我们获得了一个存放有文件路径和路径对应的标签的DataSet
对象,现在我们需要将这两个对象组合在一起。
1 | train_data = tf.data.Dataset.zip((train_file, train_label)) |
这样就组合好了,接下来调用我们上面的函数来将对应的文件载入并进行变换;由于我们将两个数据集合并到了一起,所以上面的预处理函数得稍作修改
1 | def _process_train(path,y): |
map
用于将数据集里面的每一条数据(文件路径,One-Hot格式的标签)输入到_process_train
函数中,可以使用num_parallel_calls
来指定并行单元数量来加快数据预处理
shuffle
用于再次打乱训练数据集
repeat
用于循环读取数据
batch
用于指定批次大小。如果显存大的话这里可以调大一些来加快训练速度
prefetch
用来预先读取数据
使用同样的方式构建测试数据集的流水线,但是在测试数据集中不要使用shuffle
打乱数据
构建模型
由于是快速实现那就直接调库了,不用手动一层一层搭建了
1 | model = keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=True, classes=5,weights=None) |
MobileNetV2
是一个轻量化模型,速度较快,准确率较好
keras
位于tf.keras
input_shape
为输入向量的形状,宽度和高度都是224
,具有三个颜色通道
include_top
为是否包含最终的判断概率输出层,这里选择是
classes
为分类数量,我们有五个分类选择5
weights
是否使用预训练的权重,这里不选择
编写训练代码
使用TensorFlow的keras可以快速开始训练简单的模型
1 | model = get_model() |
model.summary()
可以打印模型信息
model.compile
可以将模型配置为训练状态
optimizer
优化器,我们使用Adam
来优化我们的模型;其中学习率控制模型参数调整的幅度。过大的学习率会导致模型参数越过最优解,过小的学习率会导致模型收敛过慢
loss
损失函数。在不同问题中要选择不同的损失函数。在分类问题中选择CategoricalCrossentropy
(交叉熵)即可
metrics
优化指标,选择准确率即可
设置监控和检查点
我们使用TensorBoard来可视化监控训练过程。使用检查点来自动保存每一轮训练之后的最佳模型
1 | callbacks = [keras.callbacks.TensorBoard(), |
val_accuracy
代表在测试数据集上的准确率,这里设置为保存测试数据集下准确率最高的模型。每进行一轮训练就会判断一次;checkpoints
文件夹需要手动预先创建
TensorBoard的日志会被保存到logs
文件夹里
开始训练
一行代码即可开始训练
1 | model.fit(train_data, epochs=50, steps_per_epoch=1200, validation_data=test_data, |
训练过程分为不同轮次(epochs),一轮训练中会有很多步(step),每一步将会把一个批次传输到GPU内进行计算。一般来说传输相同数据,减少传输次数增加每次数据量可以提高速度
每一轮训练后将会在测试集上进行验证,validation_steps
控制验证步数
查看监控
1 | tensorboard --logdir .\logs\ |
接下来就可以在浏览器中查看日志,贴一个六个小时的训练结果(没好的显卡,模型性能也不怎么样)
模型部署
我们的模型当然是要用起来了!
对于移动端部署可以使用TensorFlow Lite
或者是一些各大厂商的终端推理框架
如果在PC端部署可以使用OnnxRuntime
部署,提供了预编译的多种语言SDK(有Java sdk可以包含在Mirai插件内)
模型的转换
基本上来说部署的第一步都是先将模型量化为tflite
的模型
1 | import tensorflow as tf |
到这一步就可以使用tflite在手机端使用了。如果要使用OnnxRuntime
的话还需要把tflite模型转换为onnx
模型
参考 https://github.com/onnx/tensorflow-onnx
需要注意的是在进行部署推理的时候,传入的数据也要遵循相同的预处理模式。即也需要除于255