Skip to content

Commit

Permalink
Update YOLOv3-tiny for TRT8 (wang-xinyu#1064)
Browse files Browse the repository at this point in the history 
* Update YOLOv3-tiny for TRT8
Changes from wang-xinyu#967

* Also apply fixes to utils

* Remove redundant macro definition in logging.g
  • Loading branch information
@Dominic-DallOsto
Dominic-DallOsto authored Aug 2, 2022
1 parent 0f944b8 commit 46dccf1
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 53 deletions.
3 changes: 2 additions & 1 deletion yolov3-tiny/logging.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <ostream>
#include <sstream>
#include <string>
#include "macros.h"

using Severity = nvinfer1::ILogger::Severity;

Expand Down Expand Up @@ -236,7 +237,7 @@ class Logger : public nvinfer1::ILogger
//! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
//! inheritance from nvinfer1::ILogger
//!
void log(Severity severity, const char* msg) override
void log(Severity severity, const char* msg) TRT_NOEXCEPT override
{
LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
}
Expand Down
12 changes: 12 additions & 0 deletions yolov3-tiny/macros.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
#ifndef __MACROS_H
#define __MACROS_H

#if NV_TENSORRT_MAJOR >= 8
#define TRT_NOEXCEPT noexcept
#define TRT_CONST_ENQUEUE const
#else
#define TRT_NOEXCEPT
#define TRT_CONST_ENQUEUE
#endif

#endif // __MACROS_H
5 changes: 3 additions & 2 deletions yolov3-tiny/utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include <vector>
#include <algorithm>
#include <cudnn.h>
#include "macros.h"

#ifndef CUDA_CHECK

Expand Down Expand Up @@ -38,7 +39,7 @@ namespace Tn
typedef std::pair<std::string, float> Record;
std::vector<Record> mProfile;

virtual void reportLayerTime(const char* layerName, float ms)
virtual void reportLayerTime(const char* layerName, float ms) TRT_NOEXCEPT
{
auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; });
if (record == mProfile.end())
Expand All @@ -57,7 +58,7 @@ namespace Tn

Logger(Severity severity): reportableSeverity(severity) {}

void log(Severity severity, const char* msg) override
void log(Severity severity, const char* msg) TRT_NOEXCEPT override
{
// suppress messages with severity enum value greater than the reportable
if (severity > reportableSeverity) return;
Expand Down
44 changes: 22 additions & 22 deletions yolov3-tiny/yololayer.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ namespace nvinfer1
assert(d == a + length);
}

void YoloLayerPlugin::serialize(void* buffer) const
void YoloLayerPlugin::serialize(void* buffer) const TRT_NOEXCEPT
{
using namespace Tn;
char* d = static_cast<char*>(buffer), *a = d;
Expand All @@ -50,17 +50,17 @@ namespace nvinfer1
assert(d == a + getSerializationSize());
}

size_t YoloLayerPlugin::getSerializationSize() const
size_t YoloLayerPlugin::getSerializationSize() const TRT_NOEXCEPT
{
return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size();
}

int YoloLayerPlugin::initialize()
int YoloLayerPlugin::initialize() TRT_NOEXCEPT
{
return 0;
}

Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT
{
//output the result to channel
int totalsize = MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float);
Expand All @@ -69,63 +69,63 @@ namespace nvinfer1
}

// Set plugin namespace
void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace)
void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT
{
mPluginNamespace = pluginNamespace;
}

const char* YoloLayerPlugin::getPluginNamespace() const
const char* YoloLayerPlugin::getPluginNamespace() const TRT_NOEXCEPT
{
return mPluginNamespace;
}

// Return the DataType of the plugin output at the requested index
DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT
{
return DataType::kFLOAT;
}

// Return true if output tensor is broadcast across a batch.
bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT
{
return false;
}

// Return true if plugin can use input that is broadcast across batch without replication.
bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const
bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT
{
return false;
}

void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput)
void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT
{
}

// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT
{
}

// Detach the plugin object from its execution context.
void YoloLayerPlugin::detachFromContext() {}
void YoloLayerPlugin::detachFromContext() TRT_NOEXCEPT {}

const char* YoloLayerPlugin::getPluginType() const
const char* YoloLayerPlugin::getPluginType() const TRT_NOEXCEPT
{
return "YoloLayer_TRT";
}

const char* YoloLayerPlugin::getPluginVersion() const
const char* YoloLayerPlugin::getPluginVersion() const TRT_NOEXCEPT
{
return "1";
}

void YoloLayerPlugin::destroy()
void YoloLayerPlugin::destroy() TRT_NOEXCEPT
{
delete this;
}

// Clone the plugin
IPluginV2IOExt* YoloLayerPlugin::clone() const
IPluginV2IOExt* YoloLayerPlugin::clone() const TRT_NOEXCEPT
{
YoloLayerPlugin *p = new YoloLayerPlugin();
p->setPluginNamespace(mPluginNamespace);
Expand Down Expand Up @@ -205,7 +205,7 @@ namespace nvinfer1
}


int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream)
int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT
{
//assert(batchSize == 1);
//GPU
Expand All @@ -226,29 +226,29 @@ namespace nvinfer1
mFC.fields = mPluginAttributes.data();
}

const char* YoloPluginCreator::getPluginName() const
const char* YoloPluginCreator::getPluginName() const TRT_NOEXCEPT
{
return "YoloLayer_TRT";
}

const char* YoloPluginCreator::getPluginVersion() const
const char* YoloPluginCreator::getPluginVersion() const TRT_NOEXCEPT
{
return "1";
}

const PluginFieldCollection* YoloPluginCreator::getFieldNames()
const PluginFieldCollection* YoloPluginCreator::getFieldNames() TRT_NOEXCEPT
{
return &mFC;
}

IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT
{
YoloLayerPlugin* obj = new YoloLayerPlugin();
obj->setPluginNamespace(mNamespace.c_str());
return obj;
}

IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT
{
// This object will be deleted when the network is destroyed, which will
// call MishPlugin::destroy()
Expand Down
58 changes: 30 additions & 28 deletions yolov3-tiny/yololayer.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
#include <vector>
#include <string>
#include "NvInfer.h"
#include "macros.h"


namespace Yolo
{
Expand Down Expand Up @@ -53,53 +55,53 @@ namespace nvinfer1

~YoloLayerPlugin();

int getNbOutputs() const override
int getNbOutputs() const TRT_NOEXCEPT override
{
return 1;
}

Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRT_NOEXCEPT override;

int initialize() override;
int initialize() TRT_NOEXCEPT override;

virtual void terminate() override {};
virtual void terminate() TRT_NOEXCEPT override {};

virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;}
virtual size_t getWorkspaceSize(int maxBatchSize) const TRT_NOEXCEPT override { return 0;}

virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override;
virtual int enqueue(int batchSize, const void*const * inputs, void*TRT_CONST_ENQUEUE* outputs, void* workspace, cudaStream_t stream) TRT_NOEXCEPT override;

virtual size_t getSerializationSize() const override;
virtual size_t getSerializationSize() const TRT_NOEXCEPT override;

virtual void serialize(void* buffer) const override;
virtual void serialize(void* buffer) const TRT_NOEXCEPT override;

bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override {
bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const TRT_NOEXCEPT override {
return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
}

const char* getPluginType() const override;
const char* getPluginType() const TRT_NOEXCEPT override;

const char* getPluginVersion() const override;
const char* getPluginVersion() const TRT_NOEXCEPT override;

void destroy() override;
void destroy() TRT_NOEXCEPT override;

IPluginV2IOExt* clone() const override;
IPluginV2IOExt* clone() const TRT_NOEXCEPT override;

void setPluginNamespace(const char* pluginNamespace) override;
void setPluginNamespace(const char* pluginNamespace) TRT_NOEXCEPT override;

const char* getPluginNamespace() const override;
const char* getPluginNamespace() const TRT_NOEXCEPT override;

DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const TRT_NOEXCEPT override;

bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const TRT_NOEXCEPT override;

bool canBroadcastInputAcrossBatch(int inputIndex) const override;
bool canBroadcastInputAcrossBatch(int inputIndex) const TRT_NOEXCEPT override;

void attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) TRT_NOEXCEPT override;

void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) TRT_NOEXCEPT override;

void detachFromContext() override;
void detachFromContext() TRT_NOEXCEPT override;

private:
void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1);
Expand All @@ -117,22 +119,22 @@ namespace nvinfer1

~YoloPluginCreator() override = default;

const char* getPluginName() const override;
const char* getPluginName() const TRT_NOEXCEPT override;

const char* getPluginVersion() const override;
const char* getPluginVersion() const TRT_NOEXCEPT override;

const PluginFieldCollection* getFieldNames() override;
const PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;

IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) TRT_NOEXCEPT override;

IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRT_NOEXCEPT override;

void setPluginNamespace(const char* libNamespace) override
void setPluginNamespace(const char* libNamespace) TRT_NOEXCEPT override
{
mNamespace = libNamespace;
}

const char* getPluginNamespace() const override
const char* getPluginNamespace() const TRT_NOEXCEPT override
{
return mNamespace.c_str();
}
Expand Down

0 comments on commit 46dccf1

Please sign in to comment.