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

Added BatchNorm to NetworkRT (conver_layer)

Browse Source
This commit is contained in:
perseusdg
2022-01-10 22:40:31 +05:30
parent bcf0c4eab3
commit 061bc79a69
2 changed files with 35 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/NetworkRT.h
+1
View File
@@ -98,6 +98,7 @@ public:
nvinfer1::IResizeLayer* convert_layer(nvinfer1::ITensor *input, Upsample *l); nvinfer1::IResizeLayer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input,Padding *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input,Padding *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor* input,BatchNorm *l);
#if NV_TENSORRT_MAJOR > 5 && NV_TENSORRT_MAJOR < 8 #if NV_TENSORRT_MAJOR > 5 && NV_TENSORRT_MAJOR < 8
bool serialize(const char *filename); bool serialize(const char *filename);
src/NetworkRT.cpp
+34
View File
@@ -277,6 +277,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
return convert_layer(input, (DeformConv2d*) l); return convert_layer(input, (DeformConv2d*) l);
if(type == LAYER_PADDING) if(type == LAYER_PADDING)
return convert_layer(input, (Padding*) l); return convert_layer(input, (Padding*) l);
if(type == LAYER_BATCHNORM)
return convert_layer(input,(BatchNorm*) l);
std::cout<<l->getLayerName()<<"\n"; std::cout<<l->getLayerName()<<"\n";
FatalError("Layer not implemented in tensorRT"); FatalError("Layer not implemented in tensorRT");
@@ -407,6 +409,38 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
return lRT; return lRT;
} }
ILayer* NetworkRT::convert_layer(ITensor *input,BatchNorm *l){
void *bias_b, *power_b, *mean_b, *variance_b, *scales_b;
if(dtRT == DataType::kHALF) {
bias_b = l->bias16_h;
power_b = l->power16_h;
mean_b = l->mean16_h;
variance_b = l->variance16_h;
scales_b = l->scales16_h;
} else {
bias_b = l->bias_h;
power_b = l->power_h;
mean_b = l->mean_h;
variance_b = l->variance_h;
scales_b = l->scales_h;
}
Weights power{dtRT, power_b, l->outputs};
Weights shift{dtRT, mean_b, l->outputs};
Weights scale{dtRT, variance_b, l->outputs};
IScaleLayer *lRT = networkRT->addScale(*input, ScaleMode::kCHANNEL,
shift, scale, power);
checkNULL(lRT);
Weights shift2{dtRT, bias_b, l->outputs};
Weights scale2{dtRT, scales_b, l->outputs};
IScaleLayer *lRT2 = networkRT->addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL,
shift2, scale2, power);
checkNULL(lRT2);
return lRT2;
}
ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) { ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
// std::cout<<"convert Pooling\n"; // std::cout<<"convert Pooling\n";