Add ResizeLayerRT plugin

Signed-off-by: nvidia <micaelaverucchi@gmail.com>
This commit is contained in:
nvidia
2019-12-20 15:27:41 +01:00
parent df888a3457
commit 854a4c316a
11 changed files with 171 additions and 43 deletions
+1 -1
View File
@@ -25,7 +25,7 @@ int main(int argc, char *argv[]) {
signal(SIGINT, sig_handler);
char *net = "yolo3_berkeley.rt";
char *net = "yolo3.rt";
if(argc > 1)
net = argv[1];
char *input = "../demo/yolo_test.mp4";
+1
View File
@@ -30,6 +30,7 @@ using namespace nvinfer1;
#include "pluginsRT/ShortcutRT.h"
#include "pluginsRT/YoloRT.h"
#include "pluginsRT/UpsampleRT.h"
#include "pluginsRT/ResizeLayerRT.h"
//#include "pluginsRT/Int8Calibrator.h"
class PluginFactory : IPluginFactory
+1
View File
@@ -33,6 +33,7 @@ class Yolo3Detection {
public:
int classes = 0;
int num = 0;
int n_masks = 0;
float thresh = 0.3;
cv::Scalar colors[256];
+3
View File
@@ -9,6 +9,9 @@ void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size, cud
void fill(dnnType* data, int size, dnnType val, cudaStream_t stream = cudaStream_t(0));
void resizeForward( dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w,
int o_c, int o_h, int o_w, cudaStream_t stream = cudaStream_t(0));
void reorgForward( dnnType* srcData, dnnType* dstData,
int n, int c, int h, int w, int stride, cudaStream_t stream = cudaStream_t(0));
void softmaxForward(float *input, int n, int batch, int batch_offset,
+67
View File
@@ -0,0 +1,67 @@
#include<cassert>
#include "../kernels.h"
class ResizeLayerRT : public IPlugin {
public:
ResizeLayerRT(int c, int h, int w) {
o_c = c;
o_h = h;
o_w = w;
}
~ResizeLayerRT(){
}
int getNbOutputs() const override {
return 1;
}
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
return DimsCHW{o_c, o_h, o_w};
}
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
i_c = inputDims[0].d[0];
i_h = inputDims[0].d[1];
i_w = inputDims[0].d[2];
}
int initialize() override {
return 0;
}
virtual void terminate() override {
}
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
return 0;
}
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
// printf("%d %d %d %d %d %d\n", i_c, i_w, i_h, o_c, o_w, o_h);
resizeForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
reinterpret_cast<dnnType*>(outputs[0]),
batchSize, i_c, i_h, i_w, o_c, o_h, o_w, stream);
return 0;
}
virtual size_t getSerializationSize() override {
return 6*sizeof(int);
}
virtual void serialize(void* buffer) override {
char *buf = reinterpret_cast<char*>(buffer);
tk::dnn::writeBUF(buf, o_c);
tk::dnn::writeBUF(buf, o_h);
tk::dnn::writeBUF(buf, o_w);
tk::dnn::writeBUF(buf, i_c);
tk::dnn::writeBUF(buf, i_h);
tk::dnn::writeBUF(buf, i_w);
}
int i_c, i_h, i_w, o_c, o_h, o_w;
};
+40 -32
View File
@@ -15,9 +15,9 @@ using namespace nvinfer1;
// Logger for info/warning/errors
class Logger : public ILogger {
void log(Severity severity, const char* msg) override {
// #ifdef DEBUG
#ifdef DEBUG
std::cout <<"TENSORRT LOG: "<< msg << std::endl;
// #endif
#endif
}
} loggerRT;
@@ -209,8 +209,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Dense *l) {
ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
std::cout<<"convert conv2D\n";
printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm);
// std::cout<<"convert conv2D\n";
// printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm);
void *data_b, *bias_b, *power_b, *mean_b, *variance_b, *scales_b;
@@ -261,7 +261,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
Weights power{dtRT, power_b, l->outputs};
Weights shift{dtRT, mean_b, l->outputs};
Weights scale{dtRT, variance_b, l->outputs};
std::cout<<lRT->getNbOutputs()<<std::endl;
// std::cout<<lRT->getNbOutputs()<<std::endl;
IScaleLayer *lRT2 = networkRT->addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL,
shift, scale, power);
@@ -280,8 +280,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) {
}
ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
std::cout<<"convert Pooling\n";
printf("%d %d %d %d %d %d %d %d %d %d %d %d (layer)\n", l->input_dim.h, l->input_dim.w, l->output_dim.h, l->output_dim.w, l->winW, l->winH, l->strideH, l->strideW, l->paddingH, l->paddingW, l->pool_mode, tkdnnPoolingMode_t::POOLING_MAX) ;
// std::cout<<"convert Pooling\n";
PoolingType ptype;
if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX) ptype = PoolingType::kMAX;
@@ -291,29 +290,28 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
IPoolingLayer *lRT = networkRT->addPooling(*input,
ptype, DimsHW{l->winH, l->winW});
checkNULL(lRT);
lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
lRT->setStride(DimsHW{l->strideH, l->strideW});
// if (l->input_dim.h == 13 && l->output_dim.h == 13)
// {
// lRT->setPadding(DimsHW{7, 7});
// lRT->setStride(DimsHW{2, 2});
// }
// else
// {
lRT->setPadding(DimsHW{l->paddingH, l->paddingW});
lRT->setStride(DimsHW{l->strideH, l->strideW});
// }
// IResizeLayer *lRT = networkRT->addResize(*lRT->getOutput(0));
// checkNULL(lRT);
// lRT->setOutputDimensions(l->output_dim);
ITensor *t = lRT->getOutput(0);
for(int j=0; j<t->getDimensions().nbDims; j++) {
std::cout<<t->getDimensions().d[j]<<" ";
}
std::cout<<" (TensorRT)\n";
// for(int j=0; j<t->getDimensions().nbDims; j++) {
// std::cout<<t->getDimensions().d[j]<<" ";
// }
// std::cout<<" (TensorRT)\n";
return lRT;
IPlugin *plugin = new ResizeLayerRT( l->output_dim.c,l->output_dim.h,l->output_dim.w );
IPluginLayer *lRT1 = networkRT->addPlugin(&t, 1, *plugin);
checkNULL(lRT1);
// ITensor *t1 = lRT1->getOutput(0);
// for(int j=0; j<t1->getDimensions().nbDims; j++) {
// std::cout<<t1->getDimensions().d[j]<<" ";
// }
// std::cout<<" (TensorRT after resize )\n";
return lRT1;
}
ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
@@ -347,17 +345,17 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Softmax *l) {
}
ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) {
std::cout<<"convert route\n";
// std::cout<<"convert route\n";
ITensor **tens = new ITensor*[l->layers_n];
for(int i=0; i<l->layers_n; i++) {
tens[i] = tensors[l->layers[i]];
for(int j=0; j<tens[i]->getDimensions().nbDims; j++) {
std::cout<<tens[i]->getDimensions().d[j]<<" ";
}
std::cout<<"\n";
// for(int j=0; j<tens[i]->getDimensions().nbDims; j++) {
// std::cout<<tens[i]->getDimensions().d[j]<<" ";
// }
// std::cout<<"\n";
}
IConcatenationLayer *lRT = networkRT->addConcatenation(tens, l->layers_n);
@@ -502,6 +500,16 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
return r;
}
if(name.find("Pooling") == 0) {
ResizeLayerRT *r = new ResizeLayerRT(readBUF<int>(buf), //o_c
readBUF<int>(buf), //o_h
readBUF<int>(buf)); //o_w
r->i_c = readBUF<int>(buf);
r->i_h = readBUF<int>(buf);
r->i_w = readBUF<int>(buf);
return r;
}
if(name.find("Yolo") == 0) {
YoloRT *r = new YoloRT(readBUF<int>(buf), //classes
readBUF<int>(buf), //num
-1
View File
@@ -53,7 +53,6 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW,
// checkCUDNN( cudnnGetPooling2dForwardOutputDim(poolingDesc, srcTensorDesc, &n, &c, &h, &w));
//compute w and h as in darknet
int padH = paddingH == 0? winH -1 : paddingH;
int padW = paddingW == 0? winW -1 : paddingW;
+1 -1
View File
@@ -130,7 +130,7 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int net
for (i = 0; i < lw*lh; ++i){
int row = i / lw;
int col = i % lw;
for(n = 0; n < num; ++n){
for(n = 0; n < n_masks; ++n){
int obj_index = entry_index(0, n*lw*lh + i, 4, classes, input_dim, output_dim);
float objectness = predictions[obj_index];
if(objectness <= thresh) continue;
+10 -8
View File
@@ -22,7 +22,8 @@ bool Yolo3Detection::init(std::string tensor_path) {
std::cout<<(tensor_path).c_str()<<"\n";
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
if(netRT->pluginFactory->n_yolos != 3) {
if(netRT->pluginFactory->n_yolos < 2 ) {
FatalError("this is not yolo3");
}
@@ -30,13 +31,14 @@ bool Yolo3Detection::init(std::string tensor_path) {
YoloRT *yRT = netRT->pluginFactory->yolos[i];
classes = yRT->classes;
num = yRT->num;
n_masks = yRT->n_masks;
// make a yolo layer for interpret predictions
yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, ""); // yolo without input and bias
yolo[i]->mask_h = new dnnType[num];
yolo[i]->bias_h = new dnnType[num*3*2];
memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num);
memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*2);
yolo[i] = new tk::dnn::Yolo(nullptr, classes, n_masks, ""); // yolo without input and bias
yolo[i]->mask_h = new dnnType[n_masks];
yolo[i]->bias_h = new dnnType[num*n_masks*2];
memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*n_masks);
memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*n_masks*2);
yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w);
yolo[i]->classesNames = yRT->classesNames;
}
@@ -84,7 +86,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
//DO INFERENCE
dnnType *rt_out[3];
dnnType *rt_out[netRT->pluginFactory->n_yolos];
tk::dnn::dataDim_t dim = netRT->input_dim;
checkCuda(cudaMemcpyAsync(input_d, input, dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
@@ -99,7 +101,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) {
TIMER_START
// compute dets
ndets = 0;
for(int i=0; i<3; i++) {
for(int i=0; i<netRT->pluginFactory->n_yolos; i++) {
rt_out[i] = (dnnType*)netRT->buffersRT[i+1];
yolo[i]->dstData = rt_out[i];
yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh);
+46
View File
@@ -0,0 +1,46 @@
#include "kernels.h"
#include <stdio.h>
#define MIN(a,b) (((a)<(b))?(a):(b))
#define MAX(a,b) (((a)>(b))?(a):(b))
__global__ void resize_kernel( int i_N,float *x, int i_w, int i_h, int i_c,
int o_w, int o_h, int o_c, int batch, float *out)
{
int i = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x;
if(i >= i_N) return;
int out_index = i;
int out_w = i%o_w;
i = i/o_w;
int out_h = i%o_h;
i = i/o_h;
int out_c = i%o_c;
i = i/o_c;
//copying last column/last row
int in_index = ((i*i_c + MIN(out_c,i_c-1))*i_h + MIN(out_h,i_h-1))*i_w + MIN(out_w, i_w-1);
out[out_index] = x[in_index];
}
void resizeForward( dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w,
int o_c, int o_h, int o_w, cudaStream_t stream )
{
int i_size = n*i_c*i_h*i_w;
int o_size = n*o_c*o_h*o_w;
int blocks = (o_size+255)/256;
int threads = 256;
if(i_c == o_c && i_h == o_h && i_w == o_w )
{
checkCuda(cudaMemcpy(dstData, srcData, i_size*sizeof(dnnType), cudaMemcpyDeviceToDevice));
}
else
{
checkCuda(cudaMemset(dstData, 0, o_size*sizeof(dnnType)));
resize_kernel<<<blocks, threads, 0, stream>>>(o_size, srcData, i_w, i_h, i_c, o_w, o_h, o_c, n, dstData);
// printDeviceVector(i_size, srcData);
// printDeviceVector(o_size, dstData);
}
}
+1
View File
@@ -118,5 +118,6 @@ int main() {
std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out);
std::cout<<"TRT vs correct"; checkResult(out_dim, out_data2, out);
std::cout<<"CUDNN vs TRT "; checkResult(out_dim, out_data, out_data2);
return 0;
}