pytorch_classification

Introduction

The pytorch image classification function based on torchision implemented.

Recent updates

• Updated 2022.11.05

  • Newly added tensorrt c++ inference solution

• Updated in 2022.10.29, code refactoring, basic functions are basically the same.

  • Support pytorch ddp training
  • Model reasoning that supports c++ libtorch
  • Support scripts running in one click
  • Add log module

If you are used to the previous version, please see the code of the v1 version: V1 version.

Main functions:

Using pytorch to realize image classification, based on torchision, it can extend image classification networks such as densitynet, resnext, mobilenet, efficientnet, swin transformer, etc.

If useful, welcome to star

Implement functions

• Basic functions use pytorch to realize image classification
• Includes cosine learning rate adjustment with warmup
• The step learning rate of warmup is optimized
• Multi-model fusion prediction, weighting and voting fusion
• Use flask + redis to implement model cloud API deployment (tag v1)
• Model deployment of c++ libtorch
• Augmentation for prediction when using tta test (tag v1)
• Adding pytorch implementation of label smooth (tag smoothing) (tag v1)
• Add to extract features using cnn and classify using SVM, RF, MLP, KNN and other classifiers (tag v1).
• Visualize feature layer

Running environment

• python3.7
• pytorch 1.8.1
• torchvision 0.9.1
• opencv(use of libtorch cpp inference, version 3.4.6) (optional)
• libtorch cpp reasoning use (optional)

Start quickly

Dataset Form

The organization form of the dataset, refer to sample_files/imgs/listfile.txt

Training test

Modify the parameters in run.sh , and run run.sh directly.

Main modified parameters:

 OUTPUT_PATH 模型保存和log文件的路径

TRAIN_LIST 训练数据集的list文件
VAL_LIST  测试集合的list文件
model_name 默认是resnet50
lr 学习率
epochs 训练总的epoch
batch-size  batch的大小
j dataloader的num_workers的大小
num_classes 类别数

libtorch inference

The code is stored in the cpp_inference folder.

1. Use cpp_inference/traced_model/trace_model.py to export the trained model.

2. Compile the required opencv and libtorch code to cpp_inference/third_party_library

3. Compilation

 sh compile.sh

4. Executable file testing

 ./bin/imgCls imgpath