TRT8 works with almost every nerual network now!!!!(including demo3d)

This commit is contained in:
perseusdg
2021-10-28 23:35:37 +05:30
parent 8c36dd0431
commit c5e66c6bf6
37 changed files with 1042 additions and 700 deletions
+67 -29
View File
@@ -4,22 +4,26 @@ using namespace nvinfer1;
std::vector<PluginField> ShortcutRTPluginCreator::mPluginAttributes;
PluginFieldCollection ShortcutRTPluginCreator::mFC{};
ShortcutRT::ShortcutRT(tk::dnn::dataDim_t bdim, bool mul) {
bDim = bdim;
this->bc = bDim.c;
this->bh = bDim.h;
this->bw = bDim.w;
static const char* SHORTCUTRT_PLUGIN_VERSION{"1"};
static const char* SHORTCUTRT_PLUGIN_NAME{"ShortcutRT_tkDNN"};
ShortcutRT::ShortcutRT(int bc,int bh,int bw,int c,int h,int w,bool mul) {
this->bc = bc;
this->bh = bh;
this->bw = bw;
this->mul = mul;
this->c = c;
this->h = h;
this->w = w;
}
ShortcutRT::~ShortcutRT() {}
ShortcutRT::ShortcutRT(const void *data, size_t length) {
const char* buf =reinterpret_cast<const char*>(data),*bufCheck = buf;
bDim.c = readBUF<int>(buf);
bDim.h = readBUF<int>(buf);
bDim.w = readBUF<int>(buf);
bDim.l = 1;
bc = readBUF<int>(buf);
bh = readBUF<int>(buf);
bw = readBUF<int>(buf);
mul = readBUF<bool>(buf);
c = readBUF<int>(buf);
h = readBUF<int>(buf);
@@ -35,13 +39,6 @@ Dims ShortcutRT::getOutputDimensions(int index, const Dims *inputs, int nbInputD
return Dims3{inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]};
}
void ShortcutRT::configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,
DataType type, PluginFormat format, int maxBatchSize) NOEXCEPT {
c = inputDims[0].d[0];
h = inputDims[0].d[1];
w = inputDims[0].d[2];
}
int ShortcutRT::initialize() NOEXCEPT {
return 0;
}
@@ -62,7 +59,7 @@ int ShortcutRT::enqueue(int batchSize, const void *const *inputs, void *const *o
return 0;
}
#elif NV_TENSORRT_MAJOR == 7
#elif NV_TENSORRT_MAJOR <= 7
int32_t ShortcutRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace,
cudaStream_t stream) {
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
@@ -98,11 +95,11 @@ bool ShortcutRT::supportsFormat(DataType type, PluginFormat format) const NOEXCE
}
const char *ShortcutRT::getPluginType() const NOEXCEPT {
return "ShortcutRT_tkDNN";
return SHORTCUTRT_PLUGIN_NAME;
}
const char *ShortcutRT::getPluginVersion() const NOEXCEPT {
return "1";
return SHORTCUTRT_PLUGIN_VERSION;
}
void ShortcutRT::destroy() NOEXCEPT {
@@ -117,14 +114,50 @@ void ShortcutRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
mPluginNamespace = pluginNamespace;
}
IPluginV2 *ShortcutRT::clone() const NOEXCEPT {
auto *p = new ShortcutRT(bDim,mul);
IPluginV2Ext *ShortcutRT::clone() const NOEXCEPT {
auto *p = new ShortcutRT(bc,bh,bw,c,h,w,mul);
p->setPluginNamespace(mPluginNamespace.c_str());
return p;
}
void ShortcutRT::configurePlugin(const Dims *inputDims, int32_t nbInputs, const Dims *outputDims, int32_t nbOutputs,
const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
const bool *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT {
}
bool ShortcutRT::isOutputBroadcastAcrossBatch(int32_t outputIndex, const bool *inputIsBroadcasted,
int32_t nbInputs) const NOEXCEPT {
return false;
}
bool ShortcutRT::canBroadcastInputAcrossBatch(int32_t inputIndex) const NOEXCEPT {
return false;
}
void ShortcutRT::attachToContext(cudnnContext *, cublasContext *, IGpuAllocator *) NOEXCEPT {
}
void ShortcutRT::detachFromContext() NOEXCEPT {
}
DataType ShortcutRT::getOutputDataType(int32_t index, const nvinfer1::DataType *inputTypes, int32_t nbInputs) const NOEXCEPT {
return DataType::kFLOAT;
}
ShortcutRTPluginCreator::ShortcutRTPluginCreator() {
mPluginAttributes.clear();
mPluginAttributes.emplace_back(PluginField("bc", nullptr,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("bh", nullptr,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("bw", nullptr,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("mul", nullptr,PluginFieldType::kUNKNOWN,1));
mPluginAttributes.emplace_back(PluginField("c", nullptr,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("h", nullptr,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("w", nullptr,PluginFieldType::kINT32,1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
@@ -137,28 +170,33 @@ const char *ShortcutRTPluginCreator::getPluginNamespace() const NOEXCEPT {
return mPluginNamespace.c_str();
}
IPluginV2 *ShortcutRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT {
IPluginV2Ext *ShortcutRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT {
auto *pluginObj = new ShortcutRT(serialData,serialLength);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
IPluginV2 *ShortcutRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
IPluginV2Ext *ShortcutRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
const PluginField *fields = fc->fields;
//todo assert
tk::dnn::dataDim_t bdim = *(static_cast<const tk::dnn::dataDim_t *>(fields[0].data));
bool mul = *(static_cast<const bool *>(fields[1].data));
auto *pluginObj = new ShortcutRT(bdim,mul);
assert(fc->nbFields == 7);
int bc = *(static_cast<const int *>(fields[0].data));
int bh = *(static_cast<const int *>(fields[1].data));
int bw = *(static_cast<const int *>(fields[2].data));
bool mul = *(static_cast<const bool *>(fields[3].data));
int c = *(static_cast<const int *>(fields[4].data));
int h = *(static_cast<const int *>(fields[5].data));
int w = *(static_cast<const int *>(fields[6].data));
auto *pluginObj = new ShortcutRT(bc,bh,bw,c,h,w,mul);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
const char *ShortcutRTPluginCreator::getPluginName() const NOEXCEPT {
return "ShortcutRT_tkDNN";
return SHORTCUTRT_PLUGIN_NAME;
}
const char *ShortcutRTPluginCreator::getPluginVersion() const NOEXCEPT {
return "1";
return SHORTCUTRT_PLUGIN_VERSION;
}
const PluginFieldCollection *ShortcutRTPluginCreator::getFieldNames() NOEXCEPT {