Skip to content

No360201/411_classification

Repository files navigation

Introduction

​ 该存储库是一个分类网络框架,通过增加具体的网络结构代码,即可进行训练。同时,还包含了torch->onnx->tensorrt的部署过程。

Environment

torch 1.2.0

onnx 1.4.0

tensorrt 5.1.5.0

cuda 10.1

cudnn 7.6.5

​ other dependencies in requirements.txt

How to use

train

​ 设置experiment/test/config.yaml中的mode参数为train,以及一些训练相关参数,如modeldatasets

python main.py --config_path <PATH>

​ 在deploy文件夹默认生成log.txt,保存了本次训练的loss日志

test

​ 这个过程从训练过程中保存的checkpoint模型中选择在某些指标上表现最好的模型,并可以通过配置config.yaml继续将该模型以pytorch->onnx->tensorrt的路线进行一步步的转换。

​ 设置experiment/test/config.yaml中的mode参数为test,以及配置deploy项确定是否转换到onnxtensorrt

python main.py --config_path <PATH>

​ 在当前路径下默认生成data.txt,描述网络结构和测试结果。

test one model

​ 在config.yaml中配置相应的图片文件夹路径

python demo.py --config_path <PATH> --model_path <MODEL_PATH>

tensorrt推理

​ 在config.yaml中设置mode参数为trt,并设置trt文件的路径(engine_file_path)及测试图片的路径(test_image_path

python main.py --config_path <PATH>

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published