计算机视觉一个简单易上手的图像分类任务pipeline代码

Posted 2022-12-08 江户川柯壮

整理了一个针对图像分类任务的pipeline，主要用于快速开始验证现有模型结果，以及后续的自定义模型的开发。

https://github.com/jzsherlock4869/image-classification-pipeline

对于图像分类任务来说，主要可以分为【数据】、【网络结构】、【loss函数】以及【训练步骤】这几个方面，很多情况下，这些部分都较为独立，并且也都有比较常用的形式。因此，对这几个部分进行拆解，分别设置为不同文件夹，并对各个不同的结构以单独文件的形式进行分离，保证后续可以只用添加文件就可以实现自定义模型和训练。

对于数据部分，常规的分类数据集通常有两种格式：一种是：同类的图片放在一个子文件夹，不同子文件夹代表不同的类别，比如imagenet；另一种是：所有图片都放在一起，用一个csv文件标注每个image的对应的label。代码库实现了这两种标准格式，如果你的数据集是这种格式的话，可以直接改路径就可以运行。如果不是，那么可以自定义数据集。另外，数据的augmentation也可以按照已有的样例自定义。由于内部采用了动态引入模块的方案，因此，这些改动不需要再model里对应加判断，直接config的yml中指定文件名或者类名即可。

网络结构由于有了好用的timm模型库，因此预定义了一个直接调用timm的arch，可以快速验证各个模型。如果想要自定义自己的网络结构，参照这个arch的类来写即可。

在实验和研究过程中，很多时候修改模型和修改训练策略是可以解耦的，因此将常用的操作，比如训练一次，验证并计算metric，inference等功能都封装再model类型中。如果要用特殊的训练方法，可以修改model来实现。

目前还在完善中，如有错误或问题欢迎指出~ 后续可能需要加入一些已有的常规的augment方案，提点trick，各种不同的TTA等，欢迎关注&star~