mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

(tkDNN): Support Resnet-101-AP-GeM for CUDNN, TRT to be fixed

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2022-09-24 16:57:12 +02:00
parent d4f7b4ad8b
commit 56da10df64
9 changed files with 741 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+22
View File
@@ -514,6 +514,28 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
checkNULL(lRT);
return lRT;
}
else if(l->pool_mode == tkdnnPoolingMode_t::POOLING_GENERALIZED_MEAN_P)
{
auto creator = getPluginRegistry()->getPluginCreator("GeneralizedMeanPoolingPRT_tkDNN","1");
std::vector<PluginField> mPluginAttributes;
PluginFieldCollection mFC{};
mPluginAttributes.emplace_back(PluginField("i_c",&l->input_dim.c,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("i_h",&l->input_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("i_w",&l->input_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("i_n",&l->input_dim.n,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("o_c",&l->output_dim.c,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("o_h",&l->output_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("o_w",&l->output_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("o_n",&l->output_dim.n,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("p",&l->pow_param,PluginFieldType::kFLOAT32,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;
}
else
{