Deconv tensorrt

This commit is contained in:
fbagni
2019-10-30 09:40:29 +01:00
parent f9afee2f3b
commit 33f36ab204
3 changed files with 19 additions and 10 deletions
+1 -2
View File
@@ -161,6 +161,7 @@ public:
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
int kernelH, kernelW, strideH, strideW, paddingH, paddingW; int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
bool deConv;
protected: protected:
cudnnFilterDescriptor_t filterDesc; cudnnFilterDescriptor_t filterDesc;
@@ -173,8 +174,6 @@ protected:
void inferCUDNN(dnnType* srcData, bool back = false); void inferCUDNN(dnnType* srcData, bool back = false);
void* workSpace; void* workSpace;
size_t ws_sizeInBytes; size_t ws_sizeInBytes;
bool deConv;
}; };
+1
View File
@@ -75,6 +75,7 @@ public:
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Layer *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Layer *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Conv2d *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Conv2d *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeConv2d *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Activation *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Activation *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Dense *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Dense *l);
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l);
+17 -8
View File
@@ -159,7 +159,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
if(type == LAYER_DENSE) if(type == LAYER_DENSE)
return convert_layer(input, (Dense*) l); return convert_layer(input, (Dense*) l);
if(type == LAYER_CONV2D) if(type == LAYER_CONV2D || type == LAYER_DECONV2D)
return convert_layer(input, (Conv2d*) l); return convert_layer(input, (Conv2d*) l);
if(type == LAYER_POOLING) if(type == LAYER_POOLING)
return convert_layer(input, (Pooling*) l); return convert_layer(input, (Pooling*) l);
@@ -232,13 +232,22 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
else else
b = { dtRT, nullptr, 0}; //on batchnorm bias are added later b = { dtRT, nullptr, 0}; //on batchnorm bias are added later
// Add a convolution layer with 20 outputs and a 5x5 filter. ILayer *lRT = nullptr;
IConvolutionLayer *lRT = networkRT->addConvolution(*input, if(!l->deConv) {
l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b); IConvolutionLayer *lRTconv = networkRT->addConvolution(*input,
checkNULL(lRT); l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
checkNULL(lRT);
lRT->setStride(DimsHW{l->strideH, l->strideW}); lRTconv->setStride(DimsHW{l->strideH, l->strideW});
lRT->setPadding(DimsHW{l->paddingH, l->paddingW}); lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
lRT = (ILayer*) lRTconv;
} else {
IDeconvolutionLayer *lRTconv = networkRT->addDeconvolution(*input,
l->outputs, DimsHW{l->kernelH, l->kernelW}, w, b);
checkNULL(lRT);
lRTconv->setStride(DimsHW{l->strideH, l->strideW});
lRTconv->setPadding(DimsHW{l->paddingH, l->paddingW});
lRT = (ILayer*) lRTconv;
}
if(l->batchnorm) { if(l->batchnorm) {
Weights power{dtRT, power_b, l->outputs}; Weights power{dtRT, power_b, l->outputs};