/Classification-1

轻量级图像识别框架

Primary LanguagePythonMIT LicenseMIT

亮点

功能 备注
抽象 解耦业务与框架 1. 易用：新任务仅修改Config/即可。
2. 扩展性强：各模块均设统一入口。
模型 集成Timm预训练模型库 Github stars 1. 持续更新SOTA的预训练模型(600+)。
2. 轻松定制模型。
可视化 集成TensorBoardGithub stars 可视化参数、损失、训练图像、模型结构等。
部署 服务器/移动端加速 1. TorchScript✅
2. ONNX ✅
3. TensorRT✅

快速开始

  1. 下载源码，安装依赖。

    pip install    ./Package/*zip  &&
pip install -r ./Package/requirements.txt

    注：示例数据集和参数已配好

  2. 训练：执行python train.py

  3. 测试：Config/test.yaml配置权重，执行python test.py

API文档

https://bclassification.readthedocs.io/ 文档内容包含：

  • 最佳实践

  • 自定义数据集

  • 自定义主干网络

    • 方案1：Timm库
    • 方案2：自定义
    • 附：基于Timm库定制

  • 自定义 图像增广 | 损失函数 | 优化器| 模型转换

    • 在各自入口修改即可。

  • 模型部署

    • Torch->TorchScript
    • Torch->ONNX
    • ONNX ->TensorRT

框架

├── Config
│   └── *.yaml 训练参数
│   └── *.txt  数据集 
├── DataSets
│   └── preprocess.py 图像增广入口
├── Models
│   ├── Backbone/__init__.py  主干网络入口
│   ├── Optimizer/__init__.py 优化器入口
│   ├── Loss/__init__.py      损失函数入口
│   ├── Backend               模型部署入口
├── export.py
├── test.py
└── train.py

参考