Tool of training classification models
- onnx == 1.8.1
- opencv-python == 4.2.0.32
- pytorch == 1.6.0
工程中所有可配置的超参数都在 config.py 中,有五部分(暂时还没有模块化)分别是:
- dataset
- ps1: 目前使用的是imgaug提供的数据增强pipe;所以数据增强的超参数未生效(生效于 torchvision 提供的数据增强),关于数据增强的代码在 data/transform.py 中。
- ps2: 数据集目录格式为:
- train
- cls1
- cls2
- .....
- val
- cls1
- cls2
- .....
- train
- solver
- loss
- ps1: triplet loss 暂未实现
- model
- knowledge distill
- based output 关键参数: teacher & teacher_ckpt & alpha & temperature
- based feature 关键参数: dis_feature
- 自蒸馏关键参数:cs_kd & alpha & temperature
在 config.py 中配置好参数后,执行以下命令进行训练
python train.py
训练时,模型会根据配置信息在 checkepoint 下创建目录结构并命名文件。
在 demo.py 中配置完后,执行以下命令对单张图片进行推理
python demo.py
模型训练完成之后,在 config.py 中 model & data_name & loss_name 和训练时保持一致后,运行以下命令生成 .onnx 和 .pt 文件。
python export.py