thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix network-mode

Browse Source 
Fix FP32, FP16 and INT8 modes
This commit is contained in:
unknown
2022-02-08 22:14:20 -03:00
parent 0fdaa1142a
commit d9c3cbf41d
3 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
nvdsinfer_custom_impl_Yolo/nvdsinfer_yolo_engine.cpp
View File

@@ -107,7 +107,7 @@ bool NvDsInferYoloCudaEngineGet(nvinfer1::IBuilder * const builder,
}
Yolo yolo(networkInfo);
cudaEngine = yolo.createEngine (builder);
cudaEngine = yolo.createEngine (builder, builderConfig);
if (cudaEngine == nullptr)
{
std::cerr << "Failed to build CUDA engine on "

5
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -65,7 +65,7 @@ Yolo::~Yolo()
destroyNetworkUtils();
}
nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder)
nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config)
{
assert (builder);
@@ -85,8 +85,6 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder)
std::cout << "\nNOTE: letter_box is set in cfg file, make sure to set maintain-aspect-ratio=1 in config_infer file to get better accuracy\n" << std::endl;
}
nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
if (m_NetworkMode == "INT8" && !fileExists(m_Int8CalibPath)) {
assert(builder->platformHasFastInt8());
#ifdef OPENCV
@@ -123,7 +121,6 @@ nvinfer1::ICudaEngine *Yolo::createEngine (nvinfer1::IBuilder* builder)
}
delete network;
delete config;
return engine;
}

2
nvdsinfer_custom_impl_Yolo/yolo.h
View File

@@ -74,7 +74,7 @@ public:
}
NvDsInferStatus parseModel(nvinfer1::INetworkDefinition& network) override;
nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder);
nvinfer1::ICudaEngine *createEngine (nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config);
protected:
const std::string m_NetworkType;