This commit is contained in:
Marcos Luciano
2023-06-05 16:07:37 -03:00
parent 66a6754b77
commit 79d4a0a8cd
4 changed files with 9 additions and 10 deletions

View File

@@ -9,7 +9,6 @@ NVIDIA DeepStream SDK 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1 configuration for Y
### Future updates
* DeepStream tutorials
* Updated INT8 calibration
* Support for segmentation models
* Support for classification models

View File

@@ -55,7 +55,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
assert(builder);
nvinfer1::NetworkDefinitionCreationFlags flags =
(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flags);
assert(network);
@@ -64,7 +64,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
if (m_NetworkType == "onnx") {
parser = nvonnxparser::createParser(*network, *builder->getLogger());
if (!parser->parseFromFile(m_OnnxWtsFilePath.c_str(), static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING))) {
if (!parser->parseFromFile(m_OnnxWtsFilePath.c_str(), static_cast<INT>(nvinfer1::ILogger::Severity::kWARNING))) {
std::cerr << "\nCould not parse the ONNX model\n" << std::endl;
#if NV_TENSORRT_MAJOR >= 8
@@ -99,7 +99,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
if (!m_ImplicitBatch && network->getInput(0)->getDimensions().d[0] == -1) {
nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
assert(profile);
for (int32_t i = 0; i < network->getNbInputs(); ++i) {
for (INT i = 0; i < network->getNbInputs(); ++i) {
nvinfer1::ITensor* input = network->getInput(i);
nvinfer1::Dims inputDims = input->getDimensions();
nvinfer1::Dims dims = inputDims;

View File

@@ -41,6 +41,12 @@
#include "layers/pooling_layer.h"
#include "layers/reorg_layer.h"
#if NV_TENSORRT_MAJOR >= 8
#define INT int32_t
#else
#define INT int
#endif
struct NetworkInfo
{
std::string inputBlobName;

View File

@@ -38,12 +38,6 @@
} \
}
#if NV_TENSORRT_MAJOR >= 8
#define INT int32_t
#else
#define INT int
#endif
namespace {
const char* YOLOLAYER_PLUGIN_VERSION {"1"};
const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"};