pytorch_classification

簡介

基於torchision實現的pytorch圖像分類功能。

近期更新

• 2022.11.05更新

  • 新添加tensorrt c++的推理方案

• 2022.10.29更新，進行代碼重構，基本的功能基本一致。

  • 支持pytorch ddp的訓練
  • 支持c++ libtorch的模型推理
  • 支持script腳本一鍵運行
  • 添加日誌模塊

習慣之前版本的請看v1版本的代碼：V1版本。

主要功能：

利用pytorch實現圖像分類，基於torchision可以擴展使用densenet，resnext，mobilenet，efficientnet，swin transformer等圖像分類網絡

如果有用歡迎star

實現功能

• 基礎功能利用pytorch實現圖像分類
• 包含帶有warmup的cosine學習率調整
• warmup的step學習率優調整
• 多模型融合預測，加權與投票融合
• 利用flask + redis實現模型雲端api部署（tag v1）
• c++ libtorch的模型部署
• 使用tta測試時增強進行預測（tag v1）
• 添加label smooth的pytorch實現（標籤平滑）（tag v1）
• 添加使用cnn提取特徵，並使用SVM，RF，MLP，KNN等分類器進行分類（tag v1）。
• 可視化特徵層

運行環境

• python3.7
• pytorch 1.8.1
• torchvision 0.9.1
• opencv(libtorch cpp推理使用， 版本3.4.6)（可選）
• libtorch cpp推理使用（可選）

快速開始

數據集形式

數據集的組織形式，參考sample_files/imgs/listfile.txt

訓練測試

修改run.sh中的參數，直接運行run.sh即可運行

主要修改的參數：

 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

代碼存儲在cpp_inference文件夾中。

1. 利用cpp_inference/traced_model/trace_model.py將訓練好的模型導出。

2. 編譯所需的opencv和libtorch代碼到cpp_inference/third_party_library

3. 編譯

 sh compile.sh

4. 可執行文件測試

 ./bin/imgCls imgpath