Add ResizeLayerRT plugin
Signed-off-by: nvidia <micaelaverucchi@gmail.com>
This commit is contained in:
+1
-1
@@ -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";
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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];
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user