Deconv tensorrt
This commit is contained in:
@@ -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;
|
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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};
|
||||||
|
|||||||
Reference in New Issue
Block a user