Migrate tkDNN max pooling plugin creation to pluginRegistry from the default method

This commit is contained in:
perseusdg
2021-11-10 20:09:57 +05:30
parent 7dd33cd118
commit 9e328c0daa
4 changed files with 24 additions and 6 deletions
+1
View File
@@ -500,6 +500,7 @@ public:
int winH, winW;
int strideH, strideW;
int paddingH, paddingW;
int padding;
bool size;
tkdnnPoolingMode_t pool_mode;
+15 -2
View File
@@ -369,8 +369,21 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX_FIXEDSIZE)
{
IPluginV2 *plugin = new MaxPoolFixedSizeRT(l->output_dim.c, l->output_dim.h, l->output_dim.w, l->output_dim.n, l->strideH, l->strideW, l->winH, l->winH-1);
IPluginV2Layer *lRT = networkRT->addPluginV2(&input, 1, *plugin);
auto creator = getPluginRegistry()->getPluginCreator("MaxPoolingFixedSizeRT_tkDNN","1");
std::vector<PluginField> mPluginAttributes;
PluginFieldCollection mFC{};
mPluginAttributes.emplace_back(PluginField("c",&l->output_dim.c,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("h",&l->output_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("w",&l->output_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("n",&l->output_dim.n,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("strideH",&l->strideH,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("strideW",&l->strideW,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("winSize",&l->winH,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("padding",&l->padding,PluginFieldType::kINT32,1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC);
auto *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT);
return lRT;
}
+1
View File
@@ -17,6 +17,7 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW,
this->pool_mode = pool_mode;
this->paddingH = paddingH;
this->paddingW = paddingW;
this->padding = winH -1;
checkCUDNN( cudnnCreatePoolingDescriptor(&poolingDesc) );
+7 -4
View File
@@ -4,6 +4,9 @@ using namespace nvinfer1;
std::vector<PluginField> MaxPoolFixedSizeRTPluginCreator::mPluginAttributes;
PluginFieldCollection MaxPoolFixedSizeRTPluginCreator::mFC{};
static const char* MAXPOOLFIXEDSIZERT_PLUGIN_VERSION{"1"};
static const char* MAXPOOLFIXEDSIZERT_PLUGIN_NAME{"MaxPoolingFixedSizeRT_tkDNN"};
MaxPoolFixedSizeRT::MaxPoolFixedSizeRT(int c, int h, int w, int n, int strideH, int strideW, int winSize, int padding){
this->c = c;
this->h = h;
@@ -105,11 +108,11 @@ void MaxPoolFixedSizeRT::setPluginNamespace(const char *pluginNamespace) NOEXCEP
}
const char *MaxPoolFixedSizeRT::getPluginType() const NOEXCEPT {
return "MaxPoolingFixedSizeRT_tkDNN";
return MAXPOOLFIXEDSIZERT_PLUGIN_NAME;
}
const char *MaxPoolFixedSizeRT::getPluginVersion() const NOEXCEPT {
return "1";
return MAXPOOLFIXEDSIZERT_PLUGIN_VERSION;
}
IPluginV2Ext *MaxPoolFixedSizeRT::clone() const NOEXCEPT {
@@ -185,11 +188,11 @@ IPluginV2Ext *MaxPoolFixedSizeRTPluginCreator::createPlugin(const char *name, co
}
const char *MaxPoolFixedSizeRTPluginCreator::getPluginName() const NOEXCEPT {
return "MaxPoolingFixedSizeRT_tkDNN";
return MAXPOOLFIXEDSIZERT_PLUGIN_NAME;
}
const char *MaxPoolFixedSizeRTPluginCreator::getPluginVersion() const NOEXCEPT {
return "1";
return MAXPOOLFIXEDSIZERT_PLUGIN_VERSION;
}
const PluginFieldCollection *MaxPoolFixedSizeRTPluginCreator::getFieldNames() NOEXCEPT {