TRT8 works with almost every nerual network now!!!!(including demo3d)
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user