Skip to content

Commit 3d68a2d

Browse files
committed
Update to TensorRT 8
1 parent 5660636 commit 3d68a2d

22 files changed

Lines changed: 932 additions & 590 deletions

data/settings.ini

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,8 @@ max_batch = 1
3939
gpu_id = 0
4040

4141
#-----------------------------
42-
# YOLOV2
4342
# YOLOV3
44-
# YOLOV2_TINY
45-
# YOLOV3_TINY
4643
# YOLOV4
47-
# YOLOV4_TINY
4844
# YOLOV5
4945
net_type = YOLOV4
5046

src/Detector/YoloTensorRTDetector.cpp

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,8 @@ bool YoloTensorRTDetector::Init(const config_t& config)
8989
if (net_type != config.end())
9090
{
9191
std::map<std::string, tensor_rt::ModelType> dictNetType;
92-
dictNetType["YOLOV2"] = tensor_rt::YOLOV2;
9392
dictNetType["YOLOV3"] = tensor_rt::YOLOV3;
94-
dictNetType["YOLOV2_TINY"] = tensor_rt::YOLOV2_TINY;
95-
dictNetType["YOLOV3_TINY"] = tensor_rt::YOLOV3_TINY;
9693
dictNetType["YOLOV4"] = tensor_rt::YOLOV4;
97-
dictNetType["YOLOV4_TINY"] = tensor_rt::YOLOV4_TINY;
9894
dictNetType["YOLOV5"] = tensor_rt::YOLOV5;
9995

10096
auto netType = dictNetType.find(net_type->second);

src/Detector/tensorrt_yolo/calibrator.cpp

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ SOFTWARE.
2828
#include <iostream>
2929
#include <iterator>
3030
#include <random>
31+
#include <cuda_runtime.h>
32+
#include <cuda.h>
3133

3234
Int8EntropyCalibrator::Int8EntropyCalibrator(const uint32_t& batchSize, const std::string& calibImages,
3335
const std::string& calibImagesPath,
@@ -59,7 +61,7 @@ Int8EntropyCalibrator::Int8EntropyCalibrator(const uint32_t& batchSize, const st
5961

6062
Int8EntropyCalibrator::~Int8EntropyCalibrator() { NV_CUDA_CHECK(cudaFree(m_DeviceInput)); }
6163

62-
bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int /*nbBindings*/)
64+
bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int nbBindings)noexcept
6365
{
6466
if (m_ImageIndex + m_BatchSize >= m_ImageList.size()) return false;
6567

@@ -80,7 +82,7 @@ bool Int8EntropyCalibrator::getBatch(void* bindings[], const char* names[], int
8082
return true;
8183
}
8284

83-
const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length)
85+
const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length) noexcept
8486
{
8587
void* output;
8688
m_CalibrationCache.clear();
@@ -107,7 +109,7 @@ const void* Int8EntropyCalibrator::readCalibrationCache(size_t& length)
107109
return output;
108110
}
109111

110-
void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length)
112+
void Int8EntropyCalibrator::writeCalibrationCache(const void* cache, size_t length) noexcept
111113
{
112114
assert(!m_CalibTableFilePath.empty());
113115
std::ofstream output(m_CalibTableFilePath, std::ios::binary);

src/Detector/tensorrt_yolo/calibrator.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,10 @@ class Int8EntropyCalibrator : public nvinfer1::IInt8EntropyCalibrator2
3838
const std::string& inputBlobName, const std::string &s_net_type_);
3939
virtual ~Int8EntropyCalibrator();
4040

41-
int getBatchSize() const override { return m_BatchSize; }
42-
bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
43-
const void* readCalibrationCache(size_t& length) override;
44-
void writeCalibrationCache(const void* cache, size_t length) override;
41+
int getBatchSize() const noexcept override { return m_BatchSize; }
42+
bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override;
43+
const void* readCalibrationCache(size_t& length) noexcept override;
44+
void writeCalibrationCache(const void* cache, size_t length) noexcept override;
4545

4646
private:
4747
const uint32_t m_BatchSize;

src/Detector/tensorrt_yolo/chunk.cu

Lines changed: 43 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -30,28 +30,28 @@ namespace nvinfer1
3030
{
3131

3232
}
33-
int Chunk::getNbOutputs() const
33+
int Chunk::getNbOutputs() const noexcept
3434
{
3535
return 2;
3636
}
3737

38-
Dims Chunk::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
38+
Dims Chunk::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)noexcept
3939
{
4040
assert(nbInputDims == 1);
4141
assert(index == 0 || index == 1);
4242
return Dims3(inputs[0].d[0] / 2, inputs[0].d[1], inputs[0].d[2]);
4343
}
4444

45-
int Chunk::initialize()
45+
int Chunk::initialize() noexcept
4646
{
4747
return 0;
4848
}
4949

50-
void Chunk::terminate()
50+
void Chunk::terminate() noexcept
5151
{
5252
}
5353

54-
size_t Chunk::getWorkspaceSize(int maxBatchSize) const
54+
size_t Chunk::getWorkspaceSize(int maxBatchSize) const noexcept
5555
{
5656
return 0;
5757
}
@@ -60,7 +60,16 @@ namespace nvinfer1
6060
const void* const* inputs,
6161
void** outputs,
6262
void* workspace,
63-
cudaStream_t stream)
63+
cudaStream_t stream)noexcept
64+
{
65+
return enqueue(batchSize, inputs, (void* const*)outputs, workspace, stream);
66+
}
67+
68+
int Chunk::enqueue(int batchSize,
69+
const void* const* inputs,
70+
void* const* outputs,
71+
void* workspace,
72+
cudaStream_t stream) noexcept
6473
{
6574
//batch
6675
for (int b = 0; b < batchSize; ++b)
@@ -73,54 +82,54 @@ namespace nvinfer1
7382
return 0;
7483
}
7584

76-
size_t Chunk::getSerializationSize() const
85+
size_t Chunk::getSerializationSize() const noexcept
7786
{
7887
return sizeof(_n_size_split);
7988
}
8089

81-
void Chunk::serialize(void *buffer)const
90+
void Chunk::serialize(void *buffer)const noexcept
8291
{
8392
*reinterpret_cast<int*>(buffer) = _n_size_split;
8493
}
8594

86-
const char* Chunk::getPluginType()const
95+
const char* Chunk::getPluginType()const noexcept
8796
{
8897
return "CHUNK_TRT";
8998
}
90-
const char* Chunk::getPluginVersion() const
99+
const char* Chunk::getPluginVersion() const noexcept
91100
{
92101
return "1.0";
93102
}
94103

95-
void Chunk::destroy()
104+
void Chunk::destroy() noexcept
96105
{
97106
delete this;
98107
}
99108

100-
void Chunk::setPluginNamespace(const char* pluginNamespace)
109+
void Chunk::setPluginNamespace(const char* pluginNamespace) noexcept
101110
{
102111
_s_plugin_namespace = pluginNamespace;
103112
}
104113

105-
const char* Chunk::getPluginNamespace() const
114+
const char* Chunk::getPluginNamespace() const noexcept
106115
{
107116
return _s_plugin_namespace.c_str();
108117
}
109118

110119
DataType Chunk::getOutputDataType(int index,
111120
const nvinfer1::DataType* inputTypes,
112-
int nbInputs) const
121+
int nbInputs) const noexcept
113122
{
114123
assert(index == 0 || index == 1);
115124
return DataType::kFLOAT;
116125
}
117126

118-
bool Chunk::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
127+
bool Chunk::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const noexcept
119128
{
120129
return false;
121130
}
122131

123-
bool Chunk::canBroadcastInputAcrossBatch(int inputIndex) const
132+
bool Chunk::canBroadcastInputAcrossBatch(int inputIndex) const noexcept
124133
{
125134
return false;
126135
}
@@ -133,8 +142,18 @@ namespace nvinfer1
133142
}
134143
void Chunk::detachFromContext() {}
135144

145+
bool Chunk::supportsFormat(DataType type, PluginFormat format) const noexcept
146+
{
147+
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
148+
}
149+
150+
void Chunk::configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept
151+
{
152+
153+
}
154+
136155
// Clone the plugin
137-
IPluginV2IOExt* Chunk::clone() const
156+
IPluginV2* Chunk::clone() const noexcept
138157
{
139158
Chunk *p = new Chunk();
140159
p->_n_size_split = _n_size_split;
@@ -153,41 +172,41 @@ namespace nvinfer1
153172
_fc.fields = _vec_plugin_attributes.data();
154173
}
155174

156-
const char* ChunkPluginCreator::getPluginName() const
175+
const char* ChunkPluginCreator::getPluginName() const noexcept
157176
{
158177
return "CHUNK_TRT";
159178
}
160179

161-
const char* ChunkPluginCreator::getPluginVersion() const
180+
const char* ChunkPluginCreator::getPluginVersion() const noexcept
162181
{
163182
return "1.0";
164183
}
165184

166-
const PluginFieldCollection* ChunkPluginCreator::getFieldNames()
185+
const PluginFieldCollection* ChunkPluginCreator::getFieldNames()noexcept
167186
{
168187
return &_fc;
169188
}
170189

171-
IPluginV2IOExt* ChunkPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
190+
IPluginV2* ChunkPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)noexcept
172191
{
173192
Chunk* obj = new Chunk();
174193
obj->setPluginNamespace(_s_name_space.c_str());
175194
return obj;
176195
}
177196

178-
IPluginV2IOExt* ChunkPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
197+
IPluginV2* ChunkPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)noexcept
179198
{
180199
Chunk* obj = new Chunk(serialData,serialLength);
181200
obj->setPluginNamespace(_s_name_space.c_str());
182201
return obj;
183202
}
184203

185-
void ChunkPluginCreator::setPluginNamespace(const char* libNamespace)
204+
void ChunkPluginCreator::setPluginNamespace(const char* libNamespace)noexcept
186205
{
187206
_s_name_space = libNamespace;
188207
}
189208

190-
const char* ChunkPluginCreator::getPluginNamespace() const
209+
const char* ChunkPluginCreator::getPluginNamespace() const noexcept
191210
{
192211
return _s_name_space.c_str();
193212
}

src/Detector/tensorrt_yolo/chunk.h

Lines changed: 36 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -18,37 +18,44 @@
1818

1919
namespace nvinfer1
2020
{
21-
class Chunk : public IPluginV2IOExt
21+
class Chunk : public IPluginV2
2222
{
2323
public:
2424
Chunk();
2525
Chunk(const void* buffer, size_t length);
2626
~Chunk();
27-
int getNbOutputs()const override;
28-
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
29-
int initialize() override;
30-
void terminate() override;
31-
size_t getWorkspaceSize(int maxBatchSize) const override;
32-
int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)override;
33-
size_t getSerializationSize() const override;
34-
void serialize(void* buffer) const override;
35-
const char* getPluginType() const override;
36-
const char* getPluginVersion() const override;
37-
void destroy() override;
38-
void setPluginNamespace(const char* pluginNamespace) override;
39-
const char* getPluginNamespace() const override;
40-
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
41-
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
42-
bool canBroadcastInputAcrossBatch(int inputIndex) const override;
27+
int getNbOutputs()const noexcept override;
28+
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims)noexcept override;
29+
int initialize()noexcept override;
30+
void terminate()noexcept override;
31+
size_t getWorkspaceSize(int maxBatchSize) const noexcept override;
32+
// int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream)noexcept override;
33+
34+
int enqueue(int batchSize, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept;
35+
int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) noexcept;
36+
37+
size_t getSerializationSize() const noexcept override;
38+
void serialize(void* buffer) const noexcept override;
39+
const char* getPluginType() const noexcept override;
40+
const char* getPluginVersion() const noexcept override;
41+
void destroy()noexcept override;
42+
void setPluginNamespace(const char* pluginNamespace)noexcept override;
43+
const char* getPluginNamespace() const noexcept override;
44+
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept;
45+
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const noexcept;
46+
bool canBroadcastInputAcrossBatch(int inputIndex) const noexcept;
4347
void attachToContext(
44-
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
45-
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
46-
void detachFromContext() override;
47-
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int /*nbInputs*/, int /*nbOutputs*/) const override
48+
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator);
49+
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput);
50+
void detachFromContext();
51+
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const
4852
{
4953
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
5054
}
51-
IPluginV2IOExt* clone() const override;
55+
IPluginV2* clone() const noexcept override;
56+
bool supportsFormat(DataType type, PluginFormat format) const noexcept override;
57+
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) noexcept override;
58+
5259
private:
5360
std::string _s_plugin_namespace;
5461
int _n_size_split;
@@ -59,13 +66,13 @@ namespace nvinfer1
5966
public:
6067
ChunkPluginCreator();
6168
~ChunkPluginCreator() override = default;
62-
const char* getPluginName()const override;
63-
const char* getPluginVersion() const override;
64-
const PluginFieldCollection* getFieldNames() override;
65-
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
66-
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
67-
void setPluginNamespace(const char* libNamespace) override;
68-
const char* getPluginNamespace() const override;
69+
const char* getPluginName()const noexcept override;
70+
const char* getPluginVersion() const noexcept override;
71+
const PluginFieldCollection* getFieldNames()noexcept override;
72+
IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc)noexcept override;
73+
IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength)noexcept override;
74+
void setPluginNamespace(const char* libNamespace)noexcept override;
75+
const char* getPluginNamespace() const noexcept override;
6976
private:
7077
std::string _s_name_space;
7178
static PluginFieldCollection _fc;

src/Detector/tensorrt_yolo/class_detector.h

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,12 @@ namespace tensor_rt
1919
}
2020
};
2121

22-
typedef std::vector<Result> BatchResult;
22+
using BatchResult = std::vector<Result>;
2323

2424
enum ModelType
2525
{
26-
YOLOV2 = 0,
2726
YOLOV3,
28-
YOLOV2_TINY,
29-
YOLOV3_TINY,
3027
YOLOV4,
31-
YOLOV4_TINY,
3228
YOLOV5
3329
};
3430

@@ -45,9 +41,9 @@ namespace tensor_rt
4541

4642
std::string file_model_weights = "yolov4.weights";
4743

48-
float detect_thresh = 0.9f;
44+
float detect_thresh = 0.5f;
4945

50-
ModelType net_type = YOLOV3;
46+
ModelType net_type = YOLOV4;
5147

5248
Precision inference_precison = FP32;
5349

src/Detector/tensorrt_yolo/class_timer.hpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ class Timer
3131
return elapsed();
3232
}
3333
private:
34-
typedef std::chrono::high_resolution_clock clock_;
35-
typedef std::chrono::duration<double, std::milli> second_;
34+
using clock_ = std::chrono::high_resolution_clock;
35+
using second_ = std::chrono::duration<double, std::milli>;
3636
std::chrono::time_point<clock_> beg_;
3737
};

0 commit comments

Comments
 (0)