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

Added constant and zero padding with ISliceLayer,need to add them for tensorrt versions less than 8.2 using IPluginV2Ext instead of ISliceLayer since they dont seem to support reflect and zero

Browse Source
This commit is contained in:
perseusdg
2022-01-04 09:03:04 +05:30
parent ba022663f1
commit b75cecb105
8 changed files with 396 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+30
View File
@@ -456,12 +456,42 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
}
ILayer* NetworkRT::convert_layer(ITensor *input,Padding *l){
float rt_ver = float(NV_TENSORRT_MAJOR) +
float(NV_TENSORRT_MINOR)/10 +
float(NV_TENSORRT_PATCH)/100;
#if ((NV_TENSORRT_MAJOR == 8 && NV_TENSORRT_MINOR >= 2) || NV_TENSORRT_MAJOR > 8)
auto *lRT = networkRT->addSlice(*input,Dims3{0,0,0},Dims3{l->output_dim.c,l->output_dim.h,l->output_dim.w},Dims3{0,0,0});
if(l->padding_mode == PADDING_MODE_REFLECTION){
lRT->setMode(SliceMode::kREFLECT);
}else if(l->padding_mode == PADDING_MODE_CONSTANT || l->padding_mode == PADDING_MODE_ZERO){
lRT->setMode(SliceMode::kFILL);
lRT->setInput(4, reinterpret_cast<ITensor &>(l->constant));
}
checkNULL(lRT);
return lRT;
#else
//todo add PADDING_MODE_CONSTANT AND PADDING_MODE_ZERO for tensorrt versions < 8.2
if(l->padding_mode == PADDING_MODE_REFLECTION){
auto creator = getPluginRegistry()->getPluginCreator("ReflectionPaddingRT_tkDNN","1");
std::vector<PluginField> mPluginAttributes;
PluginFieldCollection mFC{};
mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,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;
#endif
}
ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {