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

2023-05-16

整理了一个针对图像分类任务的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~

本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)

【计算机视觉】一个简单易上手的图像分类任务pipeline代码 的相关文章

随机推荐