Merge of perseusdg-tensorrt8 inside tkDNN
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
@@ -0,0 +1,186 @@
|
||||
#ifndef CENTERTRACK_H
|
||||
#define CENTERTRACK_H
|
||||
|
||||
#include <opencv2/videoio.hpp>
|
||||
#include "opencv2/opencv.hpp"
|
||||
#include "kernels.h"
|
||||
#include "utils.h"
|
||||
#include "tkdnn.h"
|
||||
#include <time.h>
|
||||
#include <vector>
|
||||
#include <numeric> // std::iota
|
||||
#include <algorithm> // std::sort
|
||||
|
||||
#include "TrackingNN.h"
|
||||
|
||||
#ifdef _WIN32
|
||||
#define _USE_MATH_DEFINES
|
||||
#include <math.h>
|
||||
#endif
|
||||
|
||||
#include "kernelsThrust.h"
|
||||
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
struct detectionRes
|
||||
{
|
||||
float score;
|
||||
int cl;
|
||||
cv::Mat ct, tr, bb0, bb1;
|
||||
float dep;
|
||||
float dim[3];
|
||||
float alpha;
|
||||
float x,y,z;
|
||||
float rot_y;
|
||||
detectionRes() : ct(cv::Mat(cv::Size(1,2), CV_32F)),
|
||||
tr(cv::Mat(cv::Size(1,2), CV_32F)),
|
||||
bb0(cv::Mat(cv::Size(1,2), CV_32F)),
|
||||
bb1(cv::Mat(cv::Size(1,2), CV_32F)) { }
|
||||
~detectionRes() {
|
||||
ct.release();
|
||||
tr.release();
|
||||
bb0.release();
|
||||
bb1.release();
|
||||
}
|
||||
};
|
||||
|
||||
struct trackingRes
|
||||
{
|
||||
struct detectionRes det_res;
|
||||
int tracking_id;
|
||||
int age;
|
||||
int active;
|
||||
int color;
|
||||
};
|
||||
|
||||
class CenterTrack : public TrackingNN
|
||||
{
|
||||
public:
|
||||
tk::dnn::dataDim_t dim;
|
||||
tk::dnn::dataDim_t dim2;
|
||||
tk::dnn::dataDim_t dim_hm;
|
||||
tk::dnn::dataDim_t dim_wh;
|
||||
tk::dnn::dataDim_t dim_reg;
|
||||
tk::dnn::dataDim_t dim_track;
|
||||
tk::dnn::dataDim_t dim_dep;
|
||||
tk::dnn::dataDim_t dim_rot;
|
||||
tk::dnn::dataDim_t dim_dim;
|
||||
tk::dnn::dataDim_t dim_amodel_offset;
|
||||
|
||||
/* preprocessing */
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
float *mean_d;
|
||||
float *stddev_d;
|
||||
#else
|
||||
cv::Vec<float, 3> mean;
|
||||
cv::Vec<float, 3> stddev;
|
||||
dnnType *input;
|
||||
#endif
|
||||
float *d_ptrs;
|
||||
|
||||
std::vector<cv::Mat> inputCalibs;
|
||||
|
||||
std::vector<cv::Size> szOld;
|
||||
|
||||
cv::Mat src;
|
||||
cv::Mat dst;
|
||||
cv::Mat dst2;
|
||||
cv::Mat trans, trans2, transOut;
|
||||
|
||||
/* pre inf */
|
||||
bool iter0;
|
||||
dnnType *input_pre_inf_d;
|
||||
bool test_pre_inf = true;
|
||||
dnnType *img_d, *hm_d;
|
||||
tk::dnn::dataDim_t dim_in0;
|
||||
tk::dnn::dataDim_t dim_in1;
|
||||
dnnType *out_d;
|
||||
|
||||
|
||||
/* postprocessing */
|
||||
int K = 100;
|
||||
int width = 128;//56; // TODO
|
||||
|
||||
// pointer used in the kernels
|
||||
float *src_out;
|
||||
int *ids_out;
|
||||
|
||||
float *topk_scores;
|
||||
int *topk_inds_;
|
||||
float *topk_ys_;
|
||||
float *topk_xs_;
|
||||
int *ids_d, *ids_;
|
||||
|
||||
float *ones;
|
||||
|
||||
float *scores, *scores_d;
|
||||
int *clses, *clses_d;
|
||||
int *topk_inds_d;
|
||||
float *topk_ys_d;
|
||||
float *topk_xs_d;
|
||||
int *inttopk_xs_d, *inttopk_ys_d;
|
||||
|
||||
float *bbx0, *bby0, *bbx1, *bby1;
|
||||
float *bbx0_d, *bby0_d, *bbx1_d, *bby1_d;
|
||||
|
||||
int *intxs, *intys;
|
||||
|
||||
float *track, *dep, *rot, *dim_, *wh, *amodel_offset;
|
||||
float *track_d, *dep_d, *rot_d, *dim_d, *wh_d, *amodel_offset_d;
|
||||
|
||||
float *target_coords;
|
||||
|
||||
/* visualization */
|
||||
cv::Mat r;
|
||||
std::vector<cv::Mat> calibs;
|
||||
cv::Mat corners, pts3DHomo;
|
||||
|
||||
std::vector<std::vector<int>> faceId;
|
||||
cv::Scalar trColors[256];
|
||||
bool mode3D;
|
||||
|
||||
//processing
|
||||
struct threshold op;
|
||||
float outThresh = 0.1;
|
||||
float newThresh = 0.3;
|
||||
// float peakThreshold = 0.2;
|
||||
// float centerThreshold = 0.3; //default 0.5
|
||||
|
||||
|
||||
//detections
|
||||
std::vector<struct detectionRes> detRes;
|
||||
int countDet;
|
||||
//tracks
|
||||
std::vector<std::vector<struct trackingRes>> trRes;
|
||||
std::vector<int> countTr;
|
||||
std::vector<int> trackId;
|
||||
|
||||
|
||||
bool init_preprocessing();
|
||||
bool init_pre_inf();
|
||||
bool init_postprocessing();
|
||||
bool init_visualization(const int n_classes);
|
||||
void pre_inf(const int bi);
|
||||
void _get_additional_inputs();
|
||||
cv::Mat transform_preds_with_trans(float x1, float x2);
|
||||
void tracking(const int bi);
|
||||
|
||||
public:
|
||||
tk::dnn::Network *pre_phase_net = nullptr;
|
||||
CenterTrack() {};
|
||||
~CenterTrack() {};
|
||||
bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1,
|
||||
const float conf_thresh=0.3, const bool mode_3d=true,
|
||||
const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>());
|
||||
void preprocess(cv::Mat &frame, const int bi=0);
|
||||
void postprocess(const int bi=0,const bool mAP=false);
|
||||
void draw(std::vector<cv::Mat>& frames);
|
||||
};
|
||||
|
||||
|
||||
} // namespace dnn
|
||||
} // namespace tk
|
||||
|
||||
|
||||
#endif /*CENTERTRACK_H*/
|
||||
@@ -73,7 +73,7 @@ public:
|
||||
CenternetDetection() {};
|
||||
~CenternetDetection() {};
|
||||
|
||||
bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
|
||||
bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
|
||||
void preprocess(cv::Mat &frame, const int bi=0);
|
||||
void postprocess(const int bi=0,const bool mAP=false);
|
||||
};
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
#ifndef CENTERNETDETECTION3D_H
|
||||
#define CENTERNETDETECTION3D_H
|
||||
|
||||
#include "kernels.h"
|
||||
#include <opencv2/videoio.hpp>
|
||||
#include "opencv2/opencv.hpp"
|
||||
#include <time.h>
|
||||
#include <vector>
|
||||
#include <numeric> // std::iota
|
||||
#include <algorithm> // std::sort
|
||||
|
||||
#ifdef _WIN32
|
||||
#define _USE_MATH_DEFINES
|
||||
#include <math.h>
|
||||
#endif
|
||||
|
||||
#include "DetectionNN3D.h"
|
||||
|
||||
#include "kernelsThrust.h"
|
||||
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
class CenternetDetection3D : public DetectionNN3D
|
||||
{
|
||||
private:
|
||||
tk::dnn::dataDim_t dim;
|
||||
tk::dnn::dataDim_t dim2;
|
||||
tk::dnn::dataDim_t dim_hm;
|
||||
tk::dnn::dataDim_t dim_wh;
|
||||
tk::dnn::dataDim_t dim_reg;
|
||||
tk::dnn::dataDim_t dim_dep;
|
||||
tk::dnn::dataDim_t dim_rot;
|
||||
tk::dnn::dataDim_t dim_dim;
|
||||
|
||||
std::vector<cv::Mat> inputCalibs;
|
||||
float *topk_scores;
|
||||
int *topk_inds_;
|
||||
float *topk_ys_;
|
||||
float *topk_xs_;
|
||||
int *ids_d, *ids_;
|
||||
|
||||
float *ones;
|
||||
|
||||
float *scores, *scores_d;
|
||||
int *clses, *clses_d;
|
||||
int *topk_inds_d;
|
||||
float *topk_ys_d;
|
||||
float *topk_xs_d;
|
||||
int *inttopk_xs_d, *inttopk_ys_d;
|
||||
|
||||
float *xs, *ys;
|
||||
|
||||
float *dep, *rot, *dim_, *wh;
|
||||
float *dep_d, *rot_d, *dim_d, *wh_d;
|
||||
|
||||
float *target_coords;
|
||||
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
float *mean_d;
|
||||
float *stddev_d;
|
||||
#else
|
||||
cv::Vec<float, 3> mean;
|
||||
cv::Vec<float, 3> stddev;
|
||||
dnnType *input;
|
||||
#endif
|
||||
cv::Mat r;
|
||||
float *d_ptrs;
|
||||
|
||||
cv::Size sz_old;
|
||||
|
||||
cv::Mat src;
|
||||
cv::Mat dst;
|
||||
cv::Mat dst2;
|
||||
cv::Mat trans, trans2;
|
||||
std::vector<cv::Mat> calibs;
|
||||
|
||||
//processing
|
||||
int K = 100;
|
||||
int width = 128;//56; // TODO
|
||||
|
||||
// pointer used in the kernels
|
||||
float *srcOut;
|
||||
int *idsOut;
|
||||
|
||||
struct threshold op;
|
||||
cv::Mat corners, pts3DHomo;
|
||||
|
||||
std::vector<std::vector<int>> faceId;
|
||||
|
||||
public:
|
||||
CenternetDetection3D() {};
|
||||
~CenternetDetection3D() {};
|
||||
|
||||
bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1, const float conf_thresh=0.3, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>());
|
||||
void preprocess(cv::Mat &frame, const int bi=0);
|
||||
void postprocess(const int bi=0,const bool mAP=false);
|
||||
void draw(std::vector<cv::Mat>& frames);
|
||||
};
|
||||
|
||||
|
||||
} // namespace dnn
|
||||
} // namespace tk
|
||||
|
||||
|
||||
#endif /*CENTERNETDETECTION_H*/
|
||||
@@ -47,5 +47,8 @@ namespace tk { namespace dnn {
|
||||
std::vector<tk::dnn::Layer*> &netLayers, const std::vector<std::string>& names);
|
||||
std::vector<std::string> darknetReadNames(const std::string& names_file);
|
||||
tk::dnn::Network* darknetParser(const std::string& cfg_file, const std::string& wgs_path, const std::string& names_file);
|
||||
void loadYoloInfo(const std::string &cfg_file,int lineNo,std::vector<float> &mask,std::vector<float> &anchors,int &num,int &classes,float &nms_thresh,int &nms_kind,int &coords);
|
||||
void loadYoloInitInfo(int &channels,int &width,int &height,const std::string &cfg_file);
|
||||
std::vector<int> noYolosLine(const std::string &cfg_file);
|
||||
|
||||
}}
|
||||
|
||||
@@ -87,7 +87,7 @@ class DetectionNN {
|
||||
* @param n_batches maximum number of batches to use in inference
|
||||
* @return true if everything is correct, false otherwise.
|
||||
*/
|
||||
virtual bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3) = 0;
|
||||
virtual bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3) = 0;
|
||||
|
||||
/**
|
||||
* This method performs the whole detection of the NN.
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
#ifndef DETECTIONNN3D_H
|
||||
#define DETECTIONNN3D_H
|
||||
|
||||
#include <iostream>
|
||||
#include <signal.h>
|
||||
#include <stdlib.h>
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include <mutex>
|
||||
#include "utils.h"
|
||||
|
||||
#include <opencv2/core/core.hpp>
|
||||
#include <opencv2/highgui/highgui.hpp>
|
||||
#include <opencv2/imgproc/imgproc.hpp>
|
||||
|
||||
#include "tkdnn.h"
|
||||
|
||||
// #define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
|
||||
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
#include <opencv2/cudawarping.hpp>
|
||||
#include <opencv2/cudaarithm.hpp>
|
||||
#endif
|
||||
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
class DetectionNN3D {
|
||||
|
||||
protected:
|
||||
tk::dnn::NetworkRT *netRT = nullptr;
|
||||
dnnType *input_d;
|
||||
|
||||
std::vector<cv::Size> originalSize;
|
||||
|
||||
cv::Scalar colors[256];
|
||||
|
||||
int nBatches = 1;
|
||||
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
cv::cuda::GpuMat bgr[3];
|
||||
cv::cuda::GpuMat imagePreproc;
|
||||
#else
|
||||
cv::Mat bgr[3];
|
||||
cv::Mat imagePreproc;
|
||||
dnnType *input;
|
||||
#endif
|
||||
|
||||
/**
|
||||
* This method preprocess the image, before feeding it to the NN.
|
||||
*
|
||||
* @param frame original frame to adapt for inference.
|
||||
* @param bi batch index
|
||||
*/
|
||||
virtual void preprocess(cv::Mat &frame, const int bi=0) = 0;
|
||||
|
||||
/**
|
||||
* This method postprocess the output of the NN to obtain the correct
|
||||
* boundig boxes.
|
||||
*
|
||||
* @param bi batch index
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation
|
||||
*/
|
||||
virtual void postprocess(const int bi=0,const bool mAP=false) = 0;
|
||||
|
||||
public:
|
||||
int classes = 0;
|
||||
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
|
||||
|
||||
std::vector<tk::dnn::box3D> detected3D; /*bounding boxes in output*/
|
||||
std::vector<std::vector<tk::dnn::box3D>> batchDetected; /*bounding boxes in output*/
|
||||
std::vector<double> pre_stats, stats, post_stats, visual_stats; /*keeps track of inference times (ms)*/
|
||||
std::vector<std::string> classesNames;
|
||||
|
||||
DetectionNN3D() {};
|
||||
~DetectionNN3D(){};
|
||||
|
||||
/**
|
||||
* Method used to initialize the class, allocate memory and compute
|
||||
* needed data.
|
||||
*
|
||||
* @param tensor_path path to the rt file of the NN.
|
||||
* @param n_classes number of classes for the given dataset.
|
||||
* @param n_batches maximum number of batches to use in inference.
|
||||
* @return true if everything is correct, false otherwise.
|
||||
*/
|
||||
virtual bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1,
|
||||
const float conf_thresh=0.3, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>()) = 0;
|
||||
|
||||
/**
|
||||
* This method performs the whole detection of the NN.
|
||||
*
|
||||
* @param frames frames to run detection on.
|
||||
* @param cur_batches number of batches to use in inference.
|
||||
* @param save_times if set to true, preprocess, inference and postprocess times
|
||||
* are saved on a csv file, otherwise not.
|
||||
* @param times pointer to the output stream where to write times.
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation.
|
||||
*/
|
||||
void update(std::vector<cv::Mat>& frames, const int cur_batches=1, bool save_times=false,
|
||||
std::ofstream *times=nullptr, const bool mAP=false){
|
||||
if(save_times && times==nullptr)
|
||||
FatalError("save_times set to true, but no valid ofstream given");
|
||||
if(cur_batches > nBatches)
|
||||
FatalError("A batch size greater than nBatches cannot be used");
|
||||
|
||||
originalSize.clear();
|
||||
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi){
|
||||
if(!frames[bi].data)
|
||||
FatalError("No image data feed to detection");
|
||||
originalSize.push_back(frames[bi].size());
|
||||
preprocess(frames[bi], bi);
|
||||
}
|
||||
TKDNN_TSTOP
|
||||
pre_stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<";";
|
||||
}
|
||||
|
||||
//do inference
|
||||
tk::dnn::dataDim_t dim = netRT->input_dim;
|
||||
dim.n = cur_batches;
|
||||
{
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
TKDNN_TSTART
|
||||
netRT->infer(dim, input_d);
|
||||
TKDNN_TSTOP
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<";";
|
||||
}
|
||||
|
||||
batchDetected.clear();
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi)
|
||||
postprocess(bi, mAP);
|
||||
TKDNN_TSTOP
|
||||
post_stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<"\n";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Method to draw bounding boxes and labels on a frame.
|
||||
*
|
||||
* @param frames original frame to draw bounding box on.
|
||||
*/
|
||||
virtual void draw(std::vector<cv::Mat>& frames){};
|
||||
|
||||
};
|
||||
|
||||
}}
|
||||
|
||||
#endif /* DETECTIONNN3D_H*/
|
||||
@@ -39,7 +39,7 @@ public:
|
||||
float *getLabels() { return mLabels.data(); }
|
||||
int getBatchesRead() const { return mBatchCount; }
|
||||
int getBatchSize() const { return mBatchSize; }
|
||||
nvinfer1::DimsNCHW getDims() const { return mDims; }
|
||||
nvinfer1::Dims4 getDims() const { return mDims; }
|
||||
float* getFileBatch() { return &mFileBatch[0]; }
|
||||
float* getFileLabels() { return &mFileLabels[0]; }
|
||||
void readInListFile(const std::string& dataFilePath, std::vector<std::string>& mListIn);
|
||||
@@ -55,7 +55,7 @@ private:
|
||||
int mFileBatchPos{ 0 };
|
||||
int mImageSize{ 0 };
|
||||
|
||||
nvinfer1::DimsNCHW mDims;
|
||||
nvinfer1::Dims4 mDims;
|
||||
std::vector<float> mBatch;
|
||||
std::vector<float> mLabels;
|
||||
std::vector<float> mFileBatch;
|
||||
|
||||
@@ -30,10 +30,10 @@ public:
|
||||
Int8EntropyCalibrator(BatchStream& stream, int firstBatch, const std::string& calibTableFilePath,
|
||||
const std::string& inputBlobName, bool readCache = true);
|
||||
virtual ~Int8EntropyCalibrator() { checkCuda(cudaFree(mDeviceInput)); }
|
||||
int getBatchSize() const override { return mStream.getBatchSize(); }
|
||||
bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
|
||||
const void* readCalibrationCache(size_t& length) override;
|
||||
void writeCalibrationCache(const void* cache, size_t length) override;
|
||||
int getBatchSize() const NOEXCEPT override { return mStream.getBatchSize(); }
|
||||
bool getBatch(void* bindings[], const char* names[], int nbBindings) NOEXCEPT override;
|
||||
const void* readCalibrationCache(size_t& length) NOEXCEPT override;
|
||||
void writeCalibrationCache(const void* cache, size_t length) NOEXCEPT override;
|
||||
|
||||
private:
|
||||
BatchStream mStream;
|
||||
|
||||
+45
-2
@@ -22,6 +22,7 @@ enum layerType_t {
|
||||
LAYER_ACTIVATION_LOGISTIC,
|
||||
LAYER_FLATTEN,
|
||||
LAYER_RESHAPE,
|
||||
LAYER_RESIZE,
|
||||
LAYER_MULADD,
|
||||
LAYER_POOLING,
|
||||
LAYER_SOFTMAX,
|
||||
@@ -55,6 +56,10 @@ public:
|
||||
|
||||
int id = 0;
|
||||
bool final; //if the layer is the final one
|
||||
unsigned int n_params = 0;
|
||||
unsigned int feature_map_size = 0;
|
||||
long unsigned MACC = 0;
|
||||
|
||||
|
||||
std::string getLayerName() {
|
||||
layerType_t type = getLayerType();
|
||||
@@ -72,6 +77,7 @@ public:
|
||||
case LAYER_ACTIVATION_LOGISTIC: return "ActivationLogistic";
|
||||
case LAYER_FLATTEN: return "Flatten";
|
||||
case LAYER_RESHAPE: return "Reshape";
|
||||
case LAYER_RESIZE: return "Resize";
|
||||
case LAYER_MULADD: return "MulAdd";
|
||||
case LAYER_POOLING: return "Pooling";
|
||||
case LAYER_SOFTMAX: return "Softmax";
|
||||
@@ -226,8 +232,9 @@ class Activation : public Layer {
|
||||
public:
|
||||
int act_mode;
|
||||
float ceiling;
|
||||
float slope;
|
||||
|
||||
Activation(Network *net, int act_mode, const float ceiling=0.0);
|
||||
Activation(Network *net, int act_mode, const float ceiling=0.0, const float slope=0.1);
|
||||
virtual ~Activation();
|
||||
virtual layerType_t getLayerType() {
|
||||
if(act_mode == CUDNN_ACTIVATION_CLIPPED_RELU)
|
||||
@@ -416,6 +423,8 @@ public:
|
||||
virtual layerType_t getLayerType() { return LAYER_FLATTEN; };
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
|
||||
int c, h, w, rows, cols;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -429,9 +438,27 @@ public:
|
||||
virtual layerType_t getLayerType() { return LAYER_RESHAPE; };
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
int n,c,h,w;
|
||||
|
||||
};
|
||||
|
||||
enum ResizeMode_t { NEAREST= 0,
|
||||
LINEAR= 1};
|
||||
|
||||
/**
|
||||
Resize layer
|
||||
*/
|
||||
class Resize : public Layer {
|
||||
|
||||
public:
|
||||
Resize(Network *net, int scale_c, int scale_h, int scale_w, bool fixed=false, ResizeMode_t mode=NEAREST);
|
||||
virtual ~Resize();
|
||||
virtual layerType_t getLayerType() { return LAYER_RESIZE; };
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
|
||||
ResizeMode_t mode;
|
||||
};
|
||||
|
||||
/**
|
||||
MulAdd layer
|
||||
@@ -473,6 +500,7 @@ public:
|
||||
int winH, winW;
|
||||
int strideH, strideW;
|
||||
int paddingH, paddingW;
|
||||
int padding;
|
||||
bool size;
|
||||
tkdnnPoolingMode_t pool_mode;
|
||||
|
||||
@@ -552,14 +580,17 @@ public:
|
||||
class Shortcut : public Layer {
|
||||
|
||||
public:
|
||||
Shortcut(Network *net, Layer *backLayer);
|
||||
Shortcut(Network *net, Layer *backLayer, bool mul=false);
|
||||
virtual ~Shortcut();
|
||||
virtual layerType_t getLayerType() { return LAYER_SHORTCUT; };
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
|
||||
int c,h,w;
|
||||
|
||||
public:
|
||||
Layer *backLayer;
|
||||
bool mul = false;
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -577,6 +608,7 @@ public:
|
||||
|
||||
int stride;
|
||||
bool reverse;
|
||||
int c,h,w;
|
||||
};
|
||||
|
||||
struct box {
|
||||
@@ -595,6 +627,16 @@ struct sortable_bbox {
|
||||
int cl;
|
||||
float **probs;
|
||||
};
|
||||
struct box3D {
|
||||
int cl;
|
||||
std::vector<float> corners;
|
||||
float prob;
|
||||
|
||||
void print()
|
||||
{
|
||||
std::cout<<"\tcl: "<<cl<<"\tprob: "<<prob<<"\tshape corners: "<<corners.size()<<std::endl;
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
Yolo3 layer
|
||||
@@ -650,6 +692,7 @@ public:
|
||||
virtual layerType_t getLayerType() { return LAYER_REGION; };
|
||||
|
||||
int classes, coords, num;
|
||||
int c,h,w;
|
||||
|
||||
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
|
||||
};
|
||||
|
||||
@@ -65,7 +65,7 @@ public:
|
||||
MobilenetDetection() {};
|
||||
~MobilenetDetection() {};
|
||||
|
||||
bool init(const std::string& tensor_path, const int n_classes, const int n_batches=1, const float conf_thresh=0.3);
|
||||
bool init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches=1, const float conf_thresh=0.3);
|
||||
void preprocess(cv::Mat &frame, const int bi=0);
|
||||
void postprocess(const int bi=0,const bool mAP=false);
|
||||
};
|
||||
|
||||
@@ -50,6 +50,7 @@ public:
|
||||
bool addLayer(Layer *l);
|
||||
void print();
|
||||
const char *getNetworkRTName(const char *network_name);
|
||||
void adjustFeatureMapSizeWithShortcuts();
|
||||
|
||||
cudnnDataType_t dataType;
|
||||
cudnnTensorFormat_t tensorFormat;
|
||||
|
||||
+32
-45
@@ -7,47 +7,28 @@
|
||||
#include "Layer.h"
|
||||
#include "NvInfer.h"
|
||||
#include <memory>
|
||||
#include <tkDNN/kernels.h>
|
||||
#include <pluginsRT/ActivationLeakyRT.h>
|
||||
#include <pluginsRT/ActivationLogisticRT.h>
|
||||
#include <pluginsRT/ActivationMishRT.h>
|
||||
#include <pluginsRT/ActivationReLUCeilingRT.h>
|
||||
#include <pluginsRT/DeformableConvRT.h>
|
||||
#include <pluginsRT/FlattenConcatRT.h>
|
||||
#include <pluginsRT/MaxPoolingFixedSizeRT.h>
|
||||
#include <pluginsRT/RegionRT.h>
|
||||
#include <pluginsRT/ReorgRT.h>
|
||||
#include <pluginsRT/ReshapeRT.h>
|
||||
#include <pluginsRT/ResizeLayerRT.h>
|
||||
#include <pluginsRT/RouteRT.h>
|
||||
#include <pluginsRT/ShortcutRT.h>
|
||||
#include <pluginsRT/UpsampleRT.h>
|
||||
#include <pluginsRT/YoloRT.h>
|
||||
|
||||
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
template<typename T> void writeBUF(char*& buffer, const T& val)
|
||||
{
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
template<typename T> T readBUF(const char*& buffer)
|
||||
{
|
||||
T val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
return val;
|
||||
}
|
||||
|
||||
using namespace nvinfer1;
|
||||
#include "pluginsRT/ActivationLeakyRT.h"
|
||||
#include "pluginsRT/ActivationLogisticRT.h"
|
||||
#include "pluginsRT/ActivationReLUCeilingRT.h"
|
||||
#include "pluginsRT/ActivationMishRT.h"
|
||||
#include "pluginsRT/ReorgRT.h"
|
||||
#include "pluginsRT/RegionRT.h"
|
||||
#include "pluginsRT/RouteRT.h"
|
||||
#include "pluginsRT/ShortcutRT.h"
|
||||
#include "pluginsRT/YoloRT.h"
|
||||
#include "pluginsRT/UpsampleRT.h"
|
||||
#include "pluginsRT/ResizeLayerRT.h"
|
||||
#include "pluginsRT/DeformableConvRT.h"
|
||||
#include "pluginsRT/FlattenConcatRT.h"
|
||||
#include "pluginsRT/ReshapeRT.h"
|
||||
#include "pluginsRT/MaxPoolingFixedSizeRT.h"
|
||||
|
||||
class PluginFactory : IPluginFactory
|
||||
{
|
||||
public:
|
||||
YoloRT *yolos[16];
|
||||
int n_yolos;
|
||||
|
||||
virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength);
|
||||
};
|
||||
|
||||
|
||||
|
||||
@@ -69,12 +50,11 @@ public:
|
||||
void* buffersRT[MAX_BUFFERS_RT];
|
||||
dataDim_t buffersDIM[MAX_BUFFERS_RT];
|
||||
int buf_input_idx, buf_output_idx;
|
||||
|
||||
bool builderActive = false;
|
||||
dataDim_t input_dim, output_dim;
|
||||
dnnType *output;
|
||||
cudaStream_t stream;
|
||||
|
||||
PluginFactory *pluginFactory;
|
||||
|
||||
NetworkRT(Network *net, const char *name);
|
||||
virtual ~NetworkRT();
|
||||
@@ -106,17 +86,24 @@ public:
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Softmax *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Route *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Flatten *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reshape *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reorg *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Region *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Flatten *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Reshape *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Resize *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Reorg *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Region *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Yolo *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Yolo *l);
|
||||
nvinfer1::IPluginV2Layer* convert_layer(nvinfer1::ITensor *input, Upsample *l);
|
||||
nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l);
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 5 && NV_TENSORRT_MAJOR < 8
|
||||
bool serialize(const char *filename);
|
||||
#else
|
||||
bool serialize(const char *filename,nvinfer1::IHostMemory *ptr);
|
||||
#endif
|
||||
|
||||
bool deserialize(const char *filename);
|
||||
void destroy();
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -5,8 +5,8 @@
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map);
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int imgdim);
|
||||
cv::Mat vizFloat2colorMap(cv::Mat map, double min=0, double max=0, int classes=19);
|
||||
cv::Mat vizData2Mat(dnnType *dataInput, tk::dnn::dataDim_t dim, int img_h, int img_w, double min=0, double max=0, int classes=19);
|
||||
cv::Mat vizLayer2Mat(tk::dnn::Network *net, int layer, int imgdim = 1000);
|
||||
|
||||
}}
|
||||
|
||||
@@ -0,0 +1,406 @@
|
||||
#ifndef SEGMENTATIONNN_H
|
||||
#define SEGMENTATIONNN_H
|
||||
|
||||
#include <iostream>
|
||||
#include <signal.h>
|
||||
#include <stdlib.h>
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
#include <mutex>
|
||||
#include "utils.h"
|
||||
|
||||
#include <opencv2/core/core.hpp>
|
||||
#include <opencv2/highgui/highgui.hpp>
|
||||
#include <opencv2/imgproc/imgproc.hpp>
|
||||
#include <opencv2/core/hal/interface.h>
|
||||
|
||||
#include "tkdnn.h"
|
||||
#include "NetworkViz.h"
|
||||
#include "kernelsThrust.h"
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
class SegmentationNN {
|
||||
|
||||
protected:
|
||||
tk::dnn::NetworkRT *netRT = nullptr;
|
||||
int nBatches = 1;
|
||||
|
||||
std::vector<cv::Size> originalSize;
|
||||
cv::Mat bgr[3];
|
||||
dnnType *input;
|
||||
dnnType *input_d;
|
||||
float* confidences_h;
|
||||
|
||||
float * tmpInputData_d;
|
||||
float *tmpOutData_d;
|
||||
float *tmpOutData_h;
|
||||
|
||||
float *mean_d, *stddev_d;
|
||||
|
||||
cublasHandle_t cublasHandle;
|
||||
|
||||
void computeBorders(const int or_width, const int or_height, int& top, int& bottom, int& left, int&right){
|
||||
top = 0;
|
||||
bottom = 0;
|
||||
left = 0;
|
||||
right = 0;
|
||||
|
||||
if(or_height != or_width){
|
||||
if(or_height < or_width){
|
||||
top = (or_width - or_height)/2;
|
||||
bottom = or_width - top - or_height;
|
||||
}
|
||||
else{
|
||||
left = (or_height - or_width)/2;
|
||||
right = or_height - left - or_width;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* This method preprocess the image, before feeding it to the NN.
|
||||
*
|
||||
* @param frame original frame to adapt for inference.
|
||||
* @param bi batch index
|
||||
*/
|
||||
void preprocess(cv::Mat &frame, const int bi=0) {
|
||||
originalSize[bi] = frame.size();
|
||||
|
||||
frame.convertTo(frame, CV_32FC3, 1 / 255.0, 0);
|
||||
int H = frame.rows;
|
||||
int W = frame.cols;
|
||||
cv::Mat frame_cropped;
|
||||
|
||||
int top, bottom, left, right;
|
||||
computeBorders(W, H, top, bottom, left, right);
|
||||
cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
|
||||
tk::dnn::dataDim_t idim = netRT->input_dim;
|
||||
|
||||
resize(frame_cropped, frame_cropped, cv::Size(idim.w, idim.h));
|
||||
|
||||
cv::split(frame_cropped, bgr);
|
||||
for (int i = 0; i < idim.c; i++){
|
||||
int idx = i * frame_cropped.rows * frame_cropped.cols;
|
||||
int ch = idim.c-1 -i;
|
||||
memcpy((void *)&input[idx + idim.tot()*bi], (void *)bgr[ch].data, frame_cropped.rows * frame_cropped.cols * sizeof(dnnType));
|
||||
}
|
||||
|
||||
checkCuda(cudaMemcpyAsync(input_d+ idim.tot()*bi, input + idim.tot()*bi, idim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
|
||||
|
||||
normalize(input_d + idim.tot()*bi, idim.c, idim.h, idim.w, mean_d, stddev_d);
|
||||
}
|
||||
|
||||
/**
|
||||
* This method postprocess the output of the NN to obtain the correct
|
||||
* boundig boxes.
|
||||
*
|
||||
* @param bi batch index
|
||||
*/
|
||||
void postprocess(const int bi=0, bool appy_colormap = true) {
|
||||
dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
|
||||
|
||||
dataDim_t odim = netRT->output_dim;
|
||||
|
||||
matrixTranspose(cublasHandle, rt_out, tmpInputData_d, odim.c, odim.w*odim.h);
|
||||
maxElem(tmpInputData_d, tmpOutData_d, odim.c, odim.h, odim.w);
|
||||
checkCuda(cudaMemcpy(tmpOutData_h, tmpOutData_d, odim.w*odim.h * sizeof(float), cudaMemcpyDeviceToHost));
|
||||
|
||||
dataDim_t vdim = odim;
|
||||
vdim.c = 1;
|
||||
|
||||
cv::Mat colored;
|
||||
|
||||
if(appy_colormap)
|
||||
colored = vizData2Mat(tmpOutData_h, vdim, netRT->input_dim.h, netRT->input_dim.w, 0, classes, classes);
|
||||
else{
|
||||
cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h);
|
||||
colored_fp32.convertTo(colored, CV_8UC1);
|
||||
}
|
||||
|
||||
int max_dim = (originalSize[bi].width > originalSize[bi].height) ? originalSize[bi].width : originalSize[bi].height;
|
||||
resize(colored, colored, cv::Size(max_dim, max_dim));
|
||||
int top, bottom, left, right;
|
||||
computeBorders(originalSize[bi].width, originalSize[bi].height, top, bottom, left, right);
|
||||
cv::Rect roi(left,top,originalSize[bi].width, originalSize[bi].height);
|
||||
cv::Mat or_size (colored, roi);
|
||||
segmented[bi] = or_size;
|
||||
};
|
||||
|
||||
public:
|
||||
int classes = 0;
|
||||
std::vector<double> stats; /*keeps track of inference times (ms)*/
|
||||
std::vector<double> stats_pre;
|
||||
std::vector<double> stats_post;
|
||||
std::vector<std::string> classesNames;
|
||||
std::vector<cv::Mat> segmented;
|
||||
|
||||
SegmentationNN() {
|
||||
checkERROR( cublasCreate(&cublasHandle) );
|
||||
};
|
||||
~SegmentationNN(){
|
||||
checkERROR( cublasDestroy(cublasHandle) );
|
||||
};
|
||||
|
||||
/**
|
||||
* Method used to inialize the class, allocate memory and compute
|
||||
* needed data.
|
||||
*
|
||||
* @param tensor_path path to the rt file og the NN.
|
||||
* @param n_classes number of classes for the given dataset.
|
||||
* @param n_batches maximum number of batches to use in inference
|
||||
* @return true if everything is correct, false otherwise.
|
||||
*/
|
||||
bool init(const std::string& tensor_path, const int n_classes=19, const int n_batches=1){
|
||||
std::cout<<(tensor_path).c_str()<<"\n";
|
||||
if(!fileExist(tensor_path.c_str()))
|
||||
FatalError("This file do not exists" + tensor_path );
|
||||
|
||||
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
|
||||
classes = n_classes;
|
||||
nBatches = n_batches;
|
||||
|
||||
checkCuda(cudaMallocHost(&input, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
|
||||
checkCuda(cudaMalloc(&input_d, sizeof(dnnType) * netRT->input_dim.tot() * nBatches));
|
||||
|
||||
dataDim_t odim = netRT->output_dim;
|
||||
|
||||
checkCuda(cudaMallocHost(&confidences_h, sizeof(float) * odim.tot()));
|
||||
checkCuda(cudaMalloc(&tmpInputData_d, sizeof(float) * odim.tot()));
|
||||
checkCuda(cudaMalloc(&tmpOutData_d, sizeof(float) * odim.w*odim.h));
|
||||
checkCuda(cudaMallocHost(&tmpOutData_h, sizeof(float) * odim.w*odim.h));
|
||||
|
||||
segmented.resize(nBatches);
|
||||
originalSize.resize(nBatches);
|
||||
|
||||
std::vector<float> mean = {0.485, 0.456, 0.406};
|
||||
std::vector<float> stddev = {0.229, 0.224, 0.225};
|
||||
|
||||
checkCuda(cudaMalloc(&mean_d, sizeof(float) * mean.size()));
|
||||
checkCuda(cudaMalloc(&stddev_d, sizeof(float) * stddev.size()));
|
||||
|
||||
checkCuda(cudaMemcpyAsync(mean_d, mean.data(), mean.size() * sizeof(float), cudaMemcpyHostToDevice, netRT->stream));
|
||||
checkCuda(cudaMemcpyAsync(stddev_d, stddev.data(), stddev.size() * sizeof(float), cudaMemcpyHostToDevice, netRT->stream));
|
||||
return true;
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* This method performs the whole detection of the NN.
|
||||
*
|
||||
* @param frames frames to run detection on.
|
||||
* @param cur_batches number of batches to use in inference
|
||||
* @param save_times if set to true, preprocess, inference and postprocess times
|
||||
* are saved on a csv file, otherwise not.
|
||||
* @param times pointer to the output stream where to write times
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation
|
||||
*/
|
||||
void update(std::vector<cv::Mat>& frames, const int cur_batches=1, bool apply_colormap=true){
|
||||
if(cur_batches > nBatches)
|
||||
FatalError("A batch size greater than nBatches cannot be used");
|
||||
|
||||
originalSize.clear();
|
||||
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi){
|
||||
if(!frames[bi].data)
|
||||
FatalError("No image data feed to detection");
|
||||
originalSize.push_back(frames[bi].size());
|
||||
preprocess(frames[bi], bi);
|
||||
}
|
||||
TKDNN_TSTOP
|
||||
stats_pre.push_back(t_ns);
|
||||
}
|
||||
|
||||
//do inference
|
||||
tk::dnn::dataDim_t dim = netRT->input_dim;
|
||||
dim.n = cur_batches;
|
||||
{
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
TKDNN_TSTART
|
||||
netRT->infer(dim, input_d);
|
||||
TKDNN_TSTOP
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
stats.push_back(t_ns);
|
||||
}
|
||||
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi)
|
||||
postprocess(bi, apply_colormap);
|
||||
TKDNN_TSTOP
|
||||
stats_post.push_back(t_ns);
|
||||
}
|
||||
}
|
||||
|
||||
void updateOriginal(cv::Mat frame, bool apply_colormap=true){
|
||||
|
||||
std::vector<cv::Mat> splitted_frames;
|
||||
int H, W, net_H, net_W;
|
||||
int top = 0, bottom = 0, left = 0, right = 0;
|
||||
std::vector<std::pair<int,int>> pos;
|
||||
|
||||
{
|
||||
TKDNN_TSTART
|
||||
cv::Size original_size = frame.size();
|
||||
|
||||
frame.convertTo(frame, CV_32FC3, 1 / 255.0, 0);
|
||||
H = frame.rows;
|
||||
W = frame.cols;
|
||||
net_H = netRT->input_dim.h;
|
||||
net_W = netRT->input_dim.w;
|
||||
|
||||
cv::Mat frame_cropped;
|
||||
|
||||
if( H <= net_H && W <= net_W ){ // smaller size wrt network
|
||||
top = (net_H - H)/2;
|
||||
bottom = net_H - H - top ;
|
||||
left = (net_W - W)/2;
|
||||
right = net_W - W - left ;
|
||||
cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0) );
|
||||
splitted_frames.push_back(frame_cropped);
|
||||
}
|
||||
else{ //bigger size wrt network
|
||||
|
||||
|
||||
if(H < net_H || W < net_W){
|
||||
if(H < net_H){
|
||||
top = (net_H - H)/2;
|
||||
bottom = net_H - H - top ;
|
||||
}
|
||||
else{
|
||||
left = (net_W - W)/2;
|
||||
right = net_W - W - left ;
|
||||
}
|
||||
cv::copyMakeBorder(frame, frame_cropped, top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(0,0,0));
|
||||
}
|
||||
|
||||
for(int x=0; x+net_W<=W ;){
|
||||
for(int y=0; y+net_H <=H ; ){
|
||||
cv::Rect roi(x, y, net_W, net_H);
|
||||
cv::Mat image_roi = frame(roi);
|
||||
splitted_frames.push_back(image_roi);
|
||||
pos.push_back(std::make_pair(x,y));
|
||||
|
||||
y += net_H;
|
||||
if(y == H)
|
||||
break;
|
||||
if(y + net_H > H) y = H - net_H;
|
||||
}
|
||||
x += net_W;
|
||||
if(x == W)
|
||||
break;
|
||||
if(x + net_W > W) x = W - net_W;
|
||||
}
|
||||
}
|
||||
|
||||
tk::dnn::dataDim_t idim = netRT->input_dim;
|
||||
|
||||
if(splitted_frames.size()> nBatches)
|
||||
FatalError(std::to_string(splitted_frames.size()) + " min batches required");
|
||||
|
||||
for(int bi=0; bi<splitted_frames.size();++bi){
|
||||
cv::split(splitted_frames[bi], bgr);
|
||||
for (int i = 0; i < idim.c; i++){
|
||||
int idx = i * splitted_frames[bi].rows * splitted_frames[bi].cols;
|
||||
int ch = idim.c-1 -i;
|
||||
memcpy((void *)&input[idx + idim.tot()*bi], (void *)bgr[ch].data, splitted_frames[bi].rows * splitted_frames[bi].cols * sizeof(dnnType));
|
||||
}
|
||||
|
||||
checkCuda(cudaMemcpyAsync(input_d+ idim.tot()*bi, input + idim.tot()*bi, idim.tot() * sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
|
||||
normalize(input_d + idim.tot()*bi, idim.c, idim.h, idim.w, mean_d, stddev_d);
|
||||
}
|
||||
TKDNN_TSTOP
|
||||
stats_pre.push_back(t_ns);
|
||||
}
|
||||
|
||||
tk::dnn::dataDim_t dim = netRT->input_dim;
|
||||
dim.n = splitted_frames.size();
|
||||
{
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
TKDNN_TSTART
|
||||
netRT->infer(dim, input_d);
|
||||
TKDNN_TSTOP
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
stats.push_back(t_ns);
|
||||
}
|
||||
|
||||
dataDim_t odim = netRT->output_dim;
|
||||
|
||||
std::vector<cv::Mat> out_img;
|
||||
|
||||
{
|
||||
TKDNN_TSTART
|
||||
|
||||
for(int bi=0; bi<splitted_frames.size();++bi){
|
||||
|
||||
dnnType *rt_out = (dnnType *)netRT->buffersRT[1]+ netRT->buffersDIM[1].tot()*bi;
|
||||
|
||||
matrixTranspose(cublasHandle, rt_out, tmpInputData_d, odim.c, odim.w*odim.h);
|
||||
maxElem(tmpInputData_d, tmpOutData_d, odim.c, odim.h, odim.w);
|
||||
checkCuda(cudaMemcpy(tmpOutData_h, tmpOutData_d, odim.w*odim.h * sizeof(float), cudaMemcpyDeviceToHost));
|
||||
|
||||
dataDim_t vdim = odim;
|
||||
vdim.c = 1;
|
||||
|
||||
cv::Mat colored;
|
||||
|
||||
if(apply_colormap)
|
||||
colored = vizData2Mat(tmpOutData_h, vdim, netRT->input_dim.h, netRT->input_dim.w, 0, classes, classes);
|
||||
else{
|
||||
cv::Mat colored_fp32 (cv::Size(odim.w, odim.h),CV_32FC1, tmpOutData_h);
|
||||
colored_fp32.convertTo(colored, CV_8UC1);
|
||||
}
|
||||
out_img.push_back(colored);
|
||||
}
|
||||
|
||||
|
||||
cv::Mat seg(frame.size(), out_img[0].type());
|
||||
if(out_img.size() == 1)
|
||||
{
|
||||
cv::Rect roi(left, top, W, H);
|
||||
seg = out_img[0](roi);
|
||||
}
|
||||
else{
|
||||
int bi=0;
|
||||
|
||||
if(top == 0 && left == 0){
|
||||
|
||||
for(int i=0; i<out_img.size(); ++i){
|
||||
cv::Mat roi_collage = seg(cv::Rect( pos[i].first ,pos[i].second,out_img[i].cols,out_img[i].rows));
|
||||
out_img[i].copyTo(roi_collage);
|
||||
}
|
||||
}
|
||||
else{
|
||||
FatalError("Not handled case")
|
||||
}
|
||||
}
|
||||
segmented[0] = seg;
|
||||
|
||||
TKDNN_TSTOP
|
||||
stats_post.push_back(t_ns);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Method to draw boundixg boxes and labels on a frame.
|
||||
*/
|
||||
cv::Mat draw(const int cur_batches=1) {
|
||||
for(int i=0; i<cur_batches; ++i){
|
||||
|
||||
cv::imshow("segmented", segmented[i]);
|
||||
cv::resizeWindow("segmented", cv::Size(512,288));
|
||||
cv::waitKey(1);
|
||||
}
|
||||
return segmented[0];
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
}}
|
||||
|
||||
#endif /* SEGMENTATIONNN_H*/
|
||||
@@ -0,0 +1,158 @@
|
||||
#ifndef TRACKINGNN_H
|
||||
#define TRACKINGNN_H
|
||||
|
||||
#include <iostream>
|
||||
#include <signal.h>
|
||||
#include <stdlib.h>
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include <mutex>
|
||||
#include "utils.h"
|
||||
|
||||
#include <opencv2/core/core.hpp>
|
||||
#include <opencv2/highgui/highgui.hpp>
|
||||
#include <opencv2/imgproc/imgproc.hpp>
|
||||
|
||||
#include "tkdnn.h"
|
||||
|
||||
// #define OPENCV_CUDACONTRIB //if OPENCV has been compiled with CUDA and contrib.
|
||||
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
#include <opencv2/cudawarping.hpp>
|
||||
#include <opencv2/cudaarithm.hpp>
|
||||
#endif
|
||||
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
class TrackingNN {
|
||||
|
||||
protected:
|
||||
tk::dnn::NetworkRT *netRT = nullptr;
|
||||
dnnType *input_d;
|
||||
|
||||
std::vector<cv::Size> originalSize;
|
||||
|
||||
cv::Scalar colors[256];
|
||||
|
||||
int nBatches = 1;
|
||||
|
||||
#ifdef OPENCV_CUDACONTRIB
|
||||
cv::cuda::GpuMat bgr[3];
|
||||
cv::cuda::GpuMat imagePreproc;
|
||||
#else
|
||||
cv::Mat bgr[3];
|
||||
cv::Mat imagePreproc;
|
||||
dnnType *input;
|
||||
#endif
|
||||
|
||||
/**
|
||||
* This method preprocess the image, before feeding it to the NN.
|
||||
*
|
||||
* @param frame original frame to adapt for inference.
|
||||
* @param bi batch index
|
||||
*/
|
||||
virtual void preprocess(cv::Mat &frame, const int bi=0) = 0;
|
||||
|
||||
/**
|
||||
* This method postprocess the output of the NN to obtain the correct
|
||||
* boundig boxes.
|
||||
*
|
||||
* @param bi batch index
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation
|
||||
*/
|
||||
virtual void postprocess(const int bi=0,const bool mAP=false) = 0;
|
||||
|
||||
public:
|
||||
int classes = 0;
|
||||
float confThreshold = 0.3; /*threshold on the confidence of the boxes*/
|
||||
|
||||
std::vector<double> pre_stats, stats, post_stats, visual_stats; /*keeps track of inference times (ms)*/
|
||||
std::vector<std::string> classesNames;
|
||||
|
||||
TrackingNN() {};
|
||||
~TrackingNN(){};
|
||||
|
||||
/**
|
||||
* Method used to initialize the class, allocate memory and compute
|
||||
* needed data.
|
||||
*
|
||||
* @param tensor_path path to the rt file of the NN.
|
||||
* @param n_classes number of classes for the given dataset.
|
||||
* @param n_batches maximum number of batches to use in inference.
|
||||
* @return true if everything is correct, false otherwise.
|
||||
*/
|
||||
virtual bool init(const std::string& tensor_path, const int n_classes=3, const int n_batches=1,
|
||||
const float conf_thresh=0.3, const bool mode_3d=true, const std::vector<cv::Mat>& k_calibs=std::vector<cv::Mat>()) = 0;
|
||||
|
||||
/**
|
||||
* This method performs the whole detection and tracking of the NN.
|
||||
*
|
||||
* @param frames frames to run detection and trcking on.
|
||||
* @param cur_batches number of batches to use in inference.
|
||||
* @param save_times if set to true, preprocess, inference and postprocess times
|
||||
* are saved on a csv file, otherwise not.
|
||||
* @param times pointer to the output stream where to write times.
|
||||
* @param mAP set to true only if all the probabilities for a bounding
|
||||
* box are needed, as in some cases for the mAP calculation.
|
||||
*/
|
||||
void update(std::vector<cv::Mat>& frames, const int cur_batches=1, bool save_times=false,
|
||||
std::ofstream *times=nullptr, const bool mAP=false){
|
||||
if(save_times && times==nullptr)
|
||||
FatalError("save_times set to true, but no valid ofstream given");
|
||||
if(cur_batches > nBatches)
|
||||
FatalError("A batch size greater than nBatches cannot be used");
|
||||
|
||||
originalSize.clear();
|
||||
if(TKDNN_VERBOSE) printCenteredTitle(" TENSORRT detection ", '=', 30);
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi){
|
||||
if(!frames[bi].data)
|
||||
FatalError("No image data feed to detection");
|
||||
originalSize.push_back(frames[bi].size());
|
||||
preprocess(frames[bi], bi);
|
||||
}
|
||||
TKDNN_TSTOP
|
||||
pre_stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<";";
|
||||
}
|
||||
|
||||
//do inference
|
||||
tk::dnn::dataDim_t dim = netRT->input_dim;
|
||||
dim.n = cur_batches;
|
||||
{
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
TKDNN_TSTART
|
||||
netRT->infer(dim, input_d);
|
||||
TKDNN_TSTOP
|
||||
if(TKDNN_VERBOSE) dim.print();
|
||||
stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<";";
|
||||
}
|
||||
|
||||
{
|
||||
TKDNN_TSTART
|
||||
for(int bi=0; bi<cur_batches;++bi)
|
||||
postprocess(bi, mAP);
|
||||
TKDNN_TSTOP
|
||||
post_stats.push_back(t_ns);
|
||||
if(save_times) *times<<t_ns<<"\n";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Method to draw bounding boxes and labels on a frame.
|
||||
*
|
||||
* @param frames original frame to draw bounding box on.
|
||||
*/
|
||||
virtual void draw(std::vector<cv::Mat>& frames){};
|
||||
|
||||
};
|
||||
|
||||
}}
|
||||
|
||||
#endif /* TRACKINGNN_H*/
|
||||
@@ -4,9 +4,9 @@
|
||||
#include "opencv2/opencv.hpp"
|
||||
|
||||
#include "DetectionNN.h"
|
||||
#include "DarknetParser.h"
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
|
||||
namespace tk { namespace dnn {
|
||||
class Yolo3Detection : public DetectionNN
|
||||
{
|
||||
private:
|
||||
@@ -19,12 +19,13 @@ private:
|
||||
tk::dnn::Yolo* getYoloLayer(int n=0);
|
||||
|
||||
cv::Mat bgr_h;
|
||||
std::vector<int> noYolos;
|
||||
|
||||
public:
|
||||
Yolo3Detection() {};
|
||||
~Yolo3Detection() {};
|
||||
|
||||
bool init(const std::string& tensor_path, const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
|
||||
bool init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path,const int n_classes=80, const int n_batches=1, const float conf_thresh=0.3);
|
||||
void preprocess(cv::Mat &frame, const int bi=0);
|
||||
void postprocess(const int bi=0,const bool mAP=false);
|
||||
};
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
#ifndef DEMO_UTILS_H
|
||||
#define DEMO_UTILS_H
|
||||
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <stdlib.h>
|
||||
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include <opencv2/core/core.hpp>
|
||||
#include <opencv2/highgui/highgui.hpp>
|
||||
#include <opencv2/imgproc/imgproc.hpp>
|
||||
|
||||
#include <yaml-cpp/yaml.h>
|
||||
|
||||
|
||||
void readCalibrationMatrix(const std::string& path, cv::Mat& calib_mat);
|
||||
|
||||
#endif //DEMO_UTILS_H
|
||||
@@ -18,6 +18,8 @@ struct Frame
|
||||
std::string iFilename;
|
||||
std::vector<BoundingBox> gt;
|
||||
std::vector<BoundingBox> det;
|
||||
int width;
|
||||
int height;
|
||||
|
||||
void print() const;
|
||||
};
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
#include "utils.h"
|
||||
|
||||
void activationELUForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
|
||||
void activationLEAKYForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
|
||||
void activationLEAKYForward(dnnType *srcData, dnnType *dstData, int size, float slope, cudaStream_t stream = cudaStream_t(0));
|
||||
void activationReLUCeilingForward(dnnType *srcData, dnnType *dstData, int size, const float ceiling, cudaStream_t stream = cudaStream_t(0));
|
||||
void activationLOGISTICForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
|
||||
void activationSIGMOIDForward(dnnType *srcData, dnnType *dstData, int size, cudaStream_t stream = cudaStream_t(0));
|
||||
@@ -24,7 +24,7 @@ void softmaxForward(float *input, int n, int batch, int batch_offset,
|
||||
int groups, int group_offset, int stride, float temp, float *output, cudaStream_t stream = cudaStream_t(0));
|
||||
|
||||
void shortcutForward(dnnType *srcData, dnnType *dstData, int n1, int c1, int h1, int w1, int s1,
|
||||
int n2, int c2, int h2, int w2, int s2,
|
||||
int n2, int c2, int h2, int w2, int s2, bool mul,
|
||||
cudaStream_t stream = cudaStream_t(0));
|
||||
|
||||
void upsampleForward(dnnType *srcData, dnnType *dstData,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
#define KERNELSTHRUST_H
|
||||
|
||||
|
||||
#include <thrust/extrema.h>
|
||||
#include <thrust/sort.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/functional.h>
|
||||
@@ -9,6 +10,8 @@
|
||||
#include <thrust/iterator/constant_iterator.h>
|
||||
#include <thrust/gather.h>
|
||||
#include <thrust/copy.h>
|
||||
#include <thrust/device_ptr.h>
|
||||
|
||||
|
||||
#include "tkdnn.h"
|
||||
|
||||
@@ -29,11 +32,15 @@ void topk(dnnType *src_begin, int *idsrc, int K, float *topk_scores,
|
||||
int *topk_inds, float *topk_ys, float *topk_xs);
|
||||
// void sortAndTopKonDevice(dnnType *src_begin, int *idsrc, float *topk_scores, int *topk_inds, float *topk_ys, float *topk_xs, const int size, const int K, const int n_classes);
|
||||
void normalize(float *bgr, const int ch, const int h, const int w, const float *mean, const float *stddev);
|
||||
void transformDep(float *src_begin, float *src_end, float *dst_begin, float *dst_end);
|
||||
void subtractWithThreshold(dnnType *src_begin, dnnType *src_end, dnnType *src2_begin, dnnType *src_out, struct threshold op);
|
||||
void topKxyclasses(int *ids_begin, int *ids_end, const int K, const int size, const int wh, int *clses, int *xs, int *ys);
|
||||
void topKxyAddOffset(int * ids_begin, const int K, const int size, int *intxs_begin, int *intys_begin,
|
||||
float *xs_begin, float *ys_begin, dnnType *src_begin, float *src_out, int *ids_out);
|
||||
void bboxes(int * ids_begin, const int K, const int size, float *xs_begin, float *ys_begin,
|
||||
dnnType *src_begin, float *bbx0, float *bbx1, float *bby0, float *bby1, float *src_out, int *ids_out);
|
||||
void getRecordsFromTopKId(int * ids_begin, const int K, const int ch, const int size, dnnType *src_begin, float *src_out, int *ids_out);
|
||||
|
||||
void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w);
|
||||
|
||||
#endif //KERNELSTHRUST_H
|
||||
@@ -1,61 +1,46 @@
|
||||
#include<cassert>
|
||||
#include "NvInfer.h"
|
||||
#include "../kernels.h"
|
||||
#include <cassert>
|
||||
#include <vector>
|
||||
|
||||
class ActivationLeakyRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class ActivationLeakyRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationLeakyRT() {
|
||||
public:
|
||||
explicit ActivationLeakyRT(float s);
|
||||
|
||||
ActivationLeakyRT(const void *data, size_t length);
|
||||
|
||||
}
|
||||
~ActivationLeakyRT();
|
||||
|
||||
~ActivationLeakyRT(){
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
int getNbOutputs() const noexcept override {
|
||||
return 1;
|
||||
}
|
||||
void
|
||||
configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) noexcept override {
|
||||
return inputs[0];
|
||||
}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
//void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
// size = 1;
|
||||
// for(int i=0; i<outputDims[0].nbDims; i++)
|
||||
// size *= outputDims[0].d[i];
|
||||
//}
|
||||
void terminate() NOEXCEPT override {}
|
||||
|
||||
int initialize() noexcept override {
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
return 0;
|
||||
}
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, void const *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
virtual void terminate() noexcept override {
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const noexcept override {
|
||||
return 0;
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
virtual int enqueue(int32_t batchSize, void const *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) noexcept override {
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
activationLEAKYForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, stream);
|
||||
return 0;
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const noexcept override {
|
||||
return 1*sizeof(int);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) const noexcept override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
|
||||
int size;
|
||||
};
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLeakyRTPluginCreator);
|
||||
};
|
||||
@@ -1,60 +1,88 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class ActivationLogisticRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
ActivationLogisticRT() {
|
||||
class ActivationLogisticRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationLogisticRT() ;
|
||||
|
||||
ActivationLogisticRT(const void *data, size_t length) ;
|
||||
|
||||
~ActivationLogisticRT() ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
~ActivationLogisticRT(){
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return inputs[0];
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
size = 1;
|
||||
for(int i=0; i<outputDims[0].nbDims; i++)
|
||||
size *= outputDims[0].d[i];
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
int initialize() override {
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
return 0;
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
int size;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
activationLOGISTICForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, stream);
|
||||
return 0;
|
||||
}
|
||||
class ActivationLogisticRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationLogisticRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 1*sizeof(int);
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer);
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
int size;
|
||||
};
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLogisticRTPluginCreator);
|
||||
};
|
||||
@@ -1,61 +1,82 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
class ActivationMishRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class ActivationMishRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationMishRT() {
|
||||
public:
|
||||
ActivationMishRT() ;
|
||||
|
||||
~ActivationMishRT() ;
|
||||
|
||||
ActivationMishRT(const void *data, size_t length) ;
|
||||
|
||||
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
~ActivationMishRT(){
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return inputs[0];
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
size = 1;
|
||||
for(int i=0; i<outputDims[0].nbDims; i++)
|
||||
size *= outputDims[0].d[i];
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
int initialize() override {
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
return 0;
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
|
||||
activationMishForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, stream);
|
||||
return 0;
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 1*sizeof(int);
|
||||
}
|
||||
void setPluginNamespace(const char *plguinNamespace) NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
int size;
|
||||
};
|
||||
int size;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class ActivationMishRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationMishRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationMishRTPluginCreator);
|
||||
};
|
||||
@@ -1,63 +1,81 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class ActivationReLUCeiling : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class ActivationReLUCeiling : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationReLUCeiling(const float ceiling) {
|
||||
this->ceiling = ceiling;
|
||||
}
|
||||
public:
|
||||
explicit ActivationReLUCeiling(const float ceiling) ;
|
||||
|
||||
~ActivationReLUCeiling(){
|
||||
~ActivationReLUCeiling() ;
|
||||
|
||||
}
|
||||
ActivationReLUCeiling(const void *data, size_t length) ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return inputs[0];
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
size = 1;
|
||||
for(int i=0; i<outputDims[0].nbDims; i++)
|
||||
size *= outputDims[0].d[i];
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int initialize() override {
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
return 0;
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
activationReLUCeilingForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, ceiling, stream);
|
||||
return 0;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 1*sizeof(int) + 1*sizeof(float);
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, ceiling);
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
assert(buf = a + getSerializationSize());
|
||||
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
int size;
|
||||
float ceiling;
|
||||
};
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
int size;
|
||||
float ceiling;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class ActivationReLUCeilingPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationReLUCeilingPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
public:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationReLUCeilingPluginCreator);
|
||||
};
|
||||
@@ -1,196 +1,137 @@
|
||||
#ifndef _DEFORMABLECONVRT_PLUGIN_H
|
||||
#define _DEFORMABLECONVRT_PLUGIN_H
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <tkdnn.h>
|
||||
|
||||
namespace nvinfer1 {
|
||||
class DeformableConvRT : public IPluginV2Ext {
|
||||
|
||||
|
||||
class DeformableConvRT : public IPlugin {
|
||||
public:
|
||||
DeformableConvRT(int chunk_dim, int kh, int kw, int sh, int sw, int ph, int pw,
|
||||
int deformableGroup, int i_n, int i_c, int i_h, int i_w,
|
||||
int o_n, int o_c, int o_h, int o_w,std::vector<dnnType> data_H,std::vector<dnnType> bias2_H,
|
||||
std::vector<dnnType> ones_d1_h,std::vector<dnnType> ones_d2_h,std::vector<dnnType> offsetH,std::vector<dnnType> maskH,int height_ones,
|
||||
int width_ones,int dim_ones);
|
||||
|
||||
~DeformableConvRT();
|
||||
|
||||
DeformableConvRT(const void *data, size_t length) ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
|
||||
cublasStatus_t stat;
|
||||
cublasHandle_t handle{nullptr};
|
||||
int i_n, i_c, i_h, i_w;
|
||||
int o_n, o_c, o_h, o_w;
|
||||
int size;
|
||||
int chunk_dim;
|
||||
int kh, kw;
|
||||
int sh, sw;
|
||||
int ph, pw;
|
||||
int deformableGroup;
|
||||
int height_ones;
|
||||
int width_ones;
|
||||
int dim_ones;
|
||||
|
||||
public:
|
||||
DeformableConvRT(int chunk_dim, int kh, int kw, int sh, int sw, int ph, int pw,
|
||||
int deformableGroup, int i_n, int i_c, int i_h, int i_w,
|
||||
int o_n, int o_c, int o_h, int o_w,
|
||||
tk::dnn::DeformConv2d *deformable = nullptr) {
|
||||
this->chunk_dim = chunk_dim;
|
||||
this->kh = kh;
|
||||
this->kw = kw;
|
||||
this->sh = sh;
|
||||
this->sw = sw;
|
||||
this->ph = ph;
|
||||
this->pw = pw;
|
||||
this->deformableGroup = deformableGroup;
|
||||
this->i_n = i_n;
|
||||
this->i_c = i_c;
|
||||
this->i_h = i_h;
|
||||
this->i_w = i_w;
|
||||
this->o_n = o_n;
|
||||
this->o_c = o_c;
|
||||
this->o_h = o_h;
|
||||
this->o_w = o_w;
|
||||
height_ones = (i_h + 2 * ph - (1 * (kh - 1) + 1)) / sh + 1;
|
||||
width_ones = (i_w + 2 * pw - (1 * (kw - 1) + 1)) / sw + 1;
|
||||
dim_ones = i_c * kh * kw * 1 * height_ones * width_ones;
|
||||
std::cout<<i_c * o_c * kh * kw * 1<<"\n";
|
||||
checkCuda( cudaMalloc(&data_d, i_c * o_c * kh * kw * 1 * sizeof(dnnType)));
|
||||
checkCuda( cudaMalloc(&bias2_d, o_c*sizeof(dnnType)));
|
||||
checkCuda( cudaMalloc(&ones_d1, height_ones * width_ones * sizeof(dnnType)));
|
||||
checkCuda( cudaMalloc(&offset, 2*chunk_dim*sizeof(dnnType)));
|
||||
checkCuda( cudaMalloc(&mask, chunk_dim*sizeof(dnnType)));
|
||||
checkCuda( cudaMalloc(&ones_d2, dim_ones*sizeof(dnnType)));
|
||||
if(deformable != nullptr) {
|
||||
this->defRT = deformable;
|
||||
checkCuda( cudaMemcpy(data_d, deformable->data_d, sizeof(dnnType)*i_c * o_c * kh * kw * 1, cudaMemcpyDeviceToDevice) );
|
||||
checkCuda( cudaMemcpy(bias2_d, deformable->bias2_d, sizeof(dnnType)*o_c, cudaMemcpyDeviceToDevice) );
|
||||
checkCuda( cudaMemcpy(ones_d1, deformable->ones_d1, sizeof(dnnType)*height_ones*width_ones, cudaMemcpyDeviceToDevice) );
|
||||
checkCuda( cudaMemcpy(offset, deformable->offset, sizeof(dnnType)*2*chunk_dim, cudaMemcpyDeviceToDevice) );
|
||||
checkCuda( cudaMemcpy(mask, deformable->mask, sizeof(dnnType)*chunk_dim, cudaMemcpyDeviceToDevice) );
|
||||
checkCuda( cudaMemcpy(ones_d2, deformable->ones_d2, sizeof(dnnType)*dim_ones, cudaMemcpyDeviceToDevice) );
|
||||
}
|
||||
stat = cublasCreate(&handle);
|
||||
if (stat != CUBLAS_STATUS_SUCCESS)
|
||||
FatalError("CUBLAS initialization failed\n");
|
||||
}
|
||||
|
||||
~DeformableConvRT() {
|
||||
checkCuda( cudaFree(data_d) );
|
||||
checkCuda( cudaFree(bias2_d) );
|
||||
checkCuda( cudaFree(ones_d1) );
|
||||
checkCuda( cudaFree(offset) );
|
||||
checkCuda( cudaFree(mask) );
|
||||
checkCuda( cudaFree(ones_d2) );
|
||||
cublasDestroy(handle);
|
||||
}
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{defRT->output_dim.c, defRT->output_dim.h, defRT->output_dim.w};
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override { }
|
||||
|
||||
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 {
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
|
||||
|
||||
// split conv2d outputs into offset to mask
|
||||
for(int b=0; b<batchSize; b++) {
|
||||
checkCuda(cudaMemcpy(offset, output_conv + b * 3 * chunk_dim, 2*chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
||||
checkCuda(cudaMemcpy(mask, output_conv + b * 3 * chunk_dim + 2*chunk_dim, chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
||||
// kernel sigmoid
|
||||
activationSIGMOIDForward(mask, mask, chunk_dim);
|
||||
// deformable convolution
|
||||
dcnV2CudaForward(stat, handle,
|
||||
srcData, data_d,
|
||||
bias2_d, ones_d1,
|
||||
offset, mask,
|
||||
reinterpret_cast<dnnType*>(outputs[0]), ones_d2,
|
||||
kh, kw,
|
||||
sh, sw,
|
||||
ph, pw,
|
||||
1, 1,
|
||||
deformableGroup, b,
|
||||
i_n, i_c, i_h, i_w,
|
||||
o_n, o_c, o_h, o_w,
|
||||
chunk_dim);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
std::vector<dnnType> data_d_v;
|
||||
std::vector<dnnType> bias2_d_v;
|
||||
std::vector<dnnType> ones_d1_v;
|
||||
std::vector<dnnType> offset_v;
|
||||
std::vector<dnnType> mask_v;
|
||||
std::vector<dnnType> ones_d2_v;
|
||||
dnnType* data_d;
|
||||
dnnType* bias2_d;
|
||||
dnnType* ones_d1;
|
||||
dnnType* offset;
|
||||
dnnType* mask;
|
||||
dnnType* ones_d2;
|
||||
// dnnType *input_n;
|
||||
// dnnType *offset_n;
|
||||
// dnnType *mask_n;
|
||||
// dnnType *output_n;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 16 * sizeof(int) + chunk_dim * 3 * sizeof(dnnType) + (i_c * o_c * kh * kw * 1 ) * sizeof(dnnType) +
|
||||
o_c * sizeof(dnnType) + height_ones * width_ones * sizeof(dnnType) + dim_ones * sizeof(dnnType);
|
||||
}
|
||||
tk::dnn::DeformConv2d *defRT;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, chunk_dim);
|
||||
tk::dnn::writeBUF(buf, kh);
|
||||
tk::dnn::writeBUF(buf, kw);
|
||||
tk::dnn::writeBUF(buf, sh);
|
||||
tk::dnn::writeBUF(buf, sw);
|
||||
tk::dnn::writeBUF(buf, ph);
|
||||
tk::dnn::writeBUF(buf, pw);
|
||||
tk::dnn::writeBUF(buf, deformableGroup);
|
||||
tk::dnn::writeBUF(buf, i_n);
|
||||
tk::dnn::writeBUF(buf, i_c);
|
||||
tk::dnn::writeBUF(buf, i_h);
|
||||
tk::dnn::writeBUF(buf, i_w);
|
||||
tk::dnn::writeBUF(buf, o_n);
|
||||
tk::dnn::writeBUF(buf, o_c);
|
||||
tk::dnn::writeBUF(buf, o_h);
|
||||
tk::dnn::writeBUF(buf, o_w);
|
||||
dnnType *aus = new dnnType[chunk_dim*2];
|
||||
checkCuda( cudaMemcpy(aus, offset, sizeof(dnnType)*2*chunk_dim, cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i<chunk_dim*2; i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
aus = new dnnType[chunk_dim];
|
||||
checkCuda( cudaMemcpy(aus, mask, sizeof(dnnType)*chunk_dim, cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i<chunk_dim; i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
aus = new dnnType[(i_c * o_c * kh * kw * 1 )];
|
||||
checkCuda( cudaMemcpy(aus, data_d, sizeof(dnnType)*(i_c * o_c * kh * kw * 1 ), cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i<(i_c * o_c * kh * kw * 1 ); i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
aus = new dnnType[o_c];
|
||||
checkCuda( cudaMemcpy(aus, bias2_d, sizeof(dnnType)*o_c, cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i < o_c; i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
aus = new dnnType[height_ones * width_ones];
|
||||
checkCuda( cudaMemcpy(aus, ones_d1, sizeof(dnnType)*height_ones * width_ones, cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i<height_ones * width_ones; i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
aus = new dnnType[dim_ones];
|
||||
checkCuda( cudaMemcpy(aus, ones_d2, sizeof(dnnType)*dim_ones, cudaMemcpyDeviceToHost) );
|
||||
for(int i=0; i<dim_ones; i++)
|
||||
tk::dnn::writeBUF(buf, aus[i]);
|
||||
free(aus);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
cublasStatus_t stat;
|
||||
cublasHandle_t handle;
|
||||
int i_n, i_c, i_h, i_w;
|
||||
int o_n, o_c, o_h, o_w;
|
||||
int size;
|
||||
int chunk_dim;
|
||||
int kh, kw;
|
||||
int sh, sw;
|
||||
int ph, pw;
|
||||
int deformableGroup;
|
||||
int height_ones;
|
||||
int width_ones;
|
||||
int dim_ones;
|
||||
|
||||
dnnType *data_d;
|
||||
dnnType *bias2_d;
|
||||
dnnType *ones_d1;
|
||||
dnnType * offset;
|
||||
dnnType * mask;
|
||||
dnnType *ones_d2;
|
||||
// dnnType *input_n;
|
||||
// dnnType *offset_n;
|
||||
// dnnType *mask_n;
|
||||
// dnnType *output_n;
|
||||
|
||||
|
||||
tk::dnn::DeformConv2d *defRT;
|
||||
class DeformableConvRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
DeformableConvRTPluginCreator();
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
REGISTER_TENSORRT_PLUGIN(DeformableConvRTPluginCreator);
|
||||
};
|
||||
#endif
|
||||
|
||||
@@ -1,81 +1,96 @@
|
||||
#include<cassert>
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
namespace nvinfer1 {
|
||||
class FlattenConcatRT : public IPluginV2Ext {
|
||||
|
||||
class FlattenConcatRT : public IPlugin {
|
||||
public:
|
||||
FlattenConcatRT(int c,int h,int w,int rows,int cols) ;
|
||||
|
||||
public:
|
||||
FlattenConcatRT() {
|
||||
stat = cublasCreate(&handle);
|
||||
if (stat != CUBLAS_STATUS_SUCCESS) {
|
||||
printf ("CUBLAS initialization failed\n");
|
||||
return;
|
||||
}
|
||||
}
|
||||
FlattenConcatRT(const void *data, size_t length) ;
|
||||
|
||||
~FlattenConcatRT(){
|
||||
~FlattenConcatRT() ;
|
||||
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{ inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2], 1, 1};
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
assert(nbOutputs == 1 && nbInputs ==1);
|
||||
rows = inputDims[0].d[0];
|
||||
cols = inputDims[0].d[1] * inputDims[0].d[2];
|
||||
c = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2];
|
||||
h = 1;
|
||||
w = 1;
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
int initialize() override {
|
||||
return 0;
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
checkERROR(cublasDestroy(handle));
|
||||
}
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR <= 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*rows*cols*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
checkERROR( cublasSetStream(handle, stream) );
|
||||
for(int i=0; i<batchSize; i++) {
|
||||
float const alpha(1.0);
|
||||
float const beta(0.0);
|
||||
int offset = i*rows*cols;
|
||||
checkERROR( cublasSgeam( handle, CUBLAS_OP_T, CUBLAS_OP_N, rows, cols, &alpha, srcData + offset, cols, &beta, srcData + offset, rows, dstData + offset, rows ));
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 5*sizeof(int);
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a = buf;
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
tk::dnn::writeBUF(buf, rows);
|
||||
tk::dnn::writeBUF(buf, cols);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
int c, h, w;
|
||||
int rows, cols;
|
||||
cublasStatus_t stat;
|
||||
cublasHandle_t handle;
|
||||
};
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
bool supportsFormat (DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
int c, h, w;
|
||||
int rows, cols;
|
||||
cublasHandle_t handle{nullptr};
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class FlattenConcatRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
FlattenConcatRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(FlattenConcatRTPluginCreator);
|
||||
};
|
||||
@@ -1,75 +1,105 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
|
||||
class MaxPoolFixedSizeRT : public IPlugin {
|
||||
|
||||
public:
|
||||
MaxPoolFixedSizeRT(int c, int h, int w, int n, int strideH, int strideW, int winSize, int padding) {
|
||||
this->c = c;
|
||||
this->h = h;
|
||||
this->w = w;
|
||||
this->n = n;
|
||||
this->stride_H = strideH;
|
||||
this->stride_W = strideW;
|
||||
this->winSize = winSize;
|
||||
this->padding = padding;
|
||||
}
|
||||
|
||||
~MaxPoolFixedSizeRT(){
|
||||
}
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{this->c, this->h, this->w};
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
}
|
||||
|
||||
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 {
|
||||
|
||||
//std::cout<<this->n<<" "<<this->c<<" "<<this->h<<" "<<this->w<<" "<<this->stride_H<<" "<<this->stride_W<<" "<<this->winSize<<" "<<this->padding<<std::endl;
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
MaxPoolingForward(srcData, dstData, batchSize, this->c, this->h, this->w, this->stride_H, this->stride_W, this->winSize, this->padding, stream);
|
||||
return 0;
|
||||
}
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 8*sizeof(int);
|
||||
}
|
||||
namespace nvinfer1 {
|
||||
class MaxPoolFixedSizeRT : public IPluginV2Ext {
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
public:
|
||||
MaxPoolFixedSizeRT(int c, int h, int w, int n, int strideH, int strideW, int winSize, int padding) ;
|
||||
|
||||
tk::dnn::writeBUF(buf, this->c);
|
||||
tk::dnn::writeBUF(buf, this->h);
|
||||
tk::dnn::writeBUF(buf, this->w);
|
||||
tk::dnn::writeBUF(buf, this->n);
|
||||
tk::dnn::writeBUF(buf, this->stride_H);
|
||||
tk::dnn::writeBUF(buf, this->stride_W);
|
||||
tk::dnn::writeBUF(buf, this->winSize);
|
||||
tk::dnn::writeBUF(buf, this->padding);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
MaxPoolFixedSizeRT(const void *data, size_t length) ;
|
||||
|
||||
int n, c, h, w;
|
||||
int stride_H, stride_W;
|
||||
int winSize;
|
||||
int padding;
|
||||
~MaxPoolFixedSizeRT() ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR <= 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
|
||||
int n, c, h, w;
|
||||
int stride_H, stride_W;
|
||||
int winSize;
|
||||
int padding;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class MaxPoolFixedSizeRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
MaxPoolFixedSizeRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(MaxPoolFixedSizeRTPluginCreator);
|
||||
};
|
||||
|
||||
@@ -1,95 +1,110 @@
|
||||
#ifndef _REGIONRT_PLUGIN_H
|
||||
#define _REGIONRT_PLUGIN_H
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class RegionRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class RegionRT : public IPluginV2Ext {
|
||||
|
||||
public:
|
||||
RegionRT(int classes, int coords, int num) {
|
||||
public:
|
||||
RegionRT(int classes, int coords, int num,int c,int h,int w);
|
||||
|
||||
this->classes = classes;
|
||||
this->coords = coords;
|
||||
this->num = num;
|
||||
}
|
||||
~RegionRT() ;
|
||||
|
||||
~RegionRT(){
|
||||
RegionRT(const void *data, size_t length) ;
|
||||
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return inputs[0];
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
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 {
|
||||
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
|
||||
for (int b = 0; b < batchSize; ++b){
|
||||
for(int n = 0; n < num; ++n){
|
||||
int index = entry_index(b, n*w*h, 0);
|
||||
activationLOGISTICForward(srcData + index, dstData + index, 2*w*h, stream);
|
||||
|
||||
index = entry_index(b, n*w*h, coords);
|
||||
activationLOGISTICForward(srcData + index, dstData + index, w*h, stream);
|
||||
}
|
||||
}
|
||||
|
||||
//softmax start
|
||||
int index = entry_index(0, 0, coords + 1);
|
||||
softmaxForward( srcData + index, classes, batchSize*num,
|
||||
(c*h*w)/num,
|
||||
w*h, 1, w*h, 1, dstData + index, stream);
|
||||
|
||||
return 0;
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, classes);
|
||||
tk::dnn::writeBUF(buf, coords);
|
||||
tk::dnn::writeBUF(buf, num);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
int c, h, w;
|
||||
int classes, coords, num;
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
int entry_index(int batch, int location, int entry) {
|
||||
int n = location / (w*h);
|
||||
int loc = location % (w*h);
|
||||
return batch*c*h*w + n*w*h*(coords+classes+1) + entry*w*h + loc;
|
||||
}
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
int c, h, w;
|
||||
int classes, coords, num;
|
||||
|
||||
int entry_index(int batch, int location, int entry) {
|
||||
int n = location / (w * h);
|
||||
int loc = location % (w * h);
|
||||
return batch * c * h * w + n * w * h * (coords + classes + 1) + entry * w * h + loc;
|
||||
}
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class RegionRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
RegionRTPluginCreator();
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(RegionRTPluginCreator);
|
||||
};
|
||||
|
||||
#endif
|
||||
|
||||
@@ -1,64 +1,98 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
class ReorgRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class ReorgRT : public IPluginV2Ext {
|
||||
|
||||
public:
|
||||
ReorgRT(int stride) {
|
||||
this->stride = stride;
|
||||
}
|
||||
public:
|
||||
ReorgRT(int stride,int c,int h,int w);
|
||||
|
||||
~ReorgRT(){
|
||||
~ReorgRT();
|
||||
|
||||
}
|
||||
ReorgRT(const void *data, size_t length);
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{inputs[0].d[0]*stride*stride, inputs[0].d[1]/stride, inputs[0].d[2]/stride};
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
int initialize() override {
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
return 0;
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
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 {
|
||||
|
||||
reorgForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]),
|
||||
batchSize, c, h, w, stride, stream);
|
||||
return 0;
|
||||
}
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, stride);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
int c, h, w, stride;
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
int c, h, w, stride;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class ReorgRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ReorgRTPluginCreator();
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ReorgRTPluginCreator);
|
||||
};
|
||||
|
||||
|
||||
@@ -1,62 +1,102 @@
|
||||
#ifndef _RESHAPERT_PLUGIN_H
|
||||
#define _RESHAPERT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
|
||||
class ReshapeRT : public IPlugin {
|
||||
|
||||
public:
|
||||
ReshapeRT(dataDim_t new_dim) {
|
||||
n = new_dim.n;
|
||||
c = new_dim.c;
|
||||
h = new_dim.h;
|
||||
w = new_dim.w;
|
||||
}
|
||||
|
||||
~ReshapeRT(){
|
||||
|
||||
}
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{ c,h,w};
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
}
|
||||
|
||||
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 {
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
return 0;
|
||||
}
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <tkdnn.h>
|
||||
using namespace tk::dnn;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
namespace nvinfer1 {
|
||||
class ReshapeRT : public IPluginV2Ext {
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a = buf;
|
||||
tk::dnn::writeBUF(buf, n);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
public:
|
||||
ReshapeRT(int n,int c,int h,int w) ;
|
||||
|
||||
int n, c, h, w;
|
||||
ReshapeRT(const void *data, size_t length) ;
|
||||
|
||||
~ReshapeRT() ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
int n, c, h, w;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class ReshapeRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ReshapeRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ReshapeRTPluginCreator);
|
||||
};
|
||||
#endif
|
||||
@@ -1,68 +1,104 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class ResizeLayerRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
ResizeLayerRT(int c, int h, int w) {
|
||||
o_c = c;
|
||||
o_h = h;
|
||||
o_w = w;
|
||||
}
|
||||
class ResizeLayerRT : public IPluginV2Ext {
|
||||
|
||||
~ResizeLayerRT(){
|
||||
}
|
||||
public:
|
||||
ResizeLayerRT(int oc, int oh, int ow,int ic,int ih,int iw) ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
ResizeLayerRT(const void *data, size_t length) ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{o_c, o_h, o_w};
|
||||
}
|
||||
~ResizeLayerRT() ;
|
||||
|
||||
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 getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
int initialize() override {
|
||||
return 0;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
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;
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR <= 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
int i_c, i_h, i_w, o_c, o_h, o_w;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class ResizeLayerRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ResizeLayerRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
|
||||
tk::dnn::writeBUF(buf, o_c);
|
||||
tk::dnn::writeBUF(buf, o_h);
|
||||
tk::dnn::writeBUF(buf, o_w);
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
|
||||
tk::dnn::writeBUF(buf, i_c);
|
||||
tk::dnn::writeBUF(buf, i_h);
|
||||
tk::dnn::writeBUF(buf, i_w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
};
|
||||
|
||||
int i_c, i_h, i_w, o_c, o_h, o_w;
|
||||
REGISTER_TENSORRT_PLUGIN(ResizeLayerRTPluginCreator);
|
||||
};
|
||||
|
||||
|
||||
@@ -1,96 +1,90 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <vector>
|
||||
#include <NvInfer.h>
|
||||
|
||||
class RouteRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class RouteRT : public IPluginV2 {
|
||||
|
||||
/**
|
||||
THIS IS NOT USED ANYMORE
|
||||
*/
|
||||
/**
|
||||
THIS IS NOT USED ANYMORE
|
||||
*/
|
||||
|
||||
public:
|
||||
RouteRT(int groups, int group_id) {
|
||||
this->groups = groups;
|
||||
this->group_id = group_id;
|
||||
}
|
||||
public:
|
||||
RouteRT(int groups, int group_id) ;
|
||||
|
||||
~RouteRT(){
|
||||
~RouteRT() ;
|
||||
|
||||
}
|
||||
RouteRT(const void *data, size_t length) ;
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
int out_c = 0;
|
||||
for(int i=0; i<nbInputDims; i++) out_c += inputs[i].d[0];
|
||||
return DimsCHW{out_c/groups, inputs[0].d[1], inputs[0].d[2]};
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
in = nbInputs;
|
||||
c = 0;
|
||||
for(int i=0; i<nbInputs; i++) {
|
||||
c_in[i] = inputDims[i].d[0];
|
||||
c += inputDims[i].d[0];
|
||||
}
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
c /= groups;
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int initialize() override {
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
return 0;
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const override {
|
||||
return 0;
|
||||
}
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
|
||||
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
for(int b=0; b<batchSize; b++) {
|
||||
int offset = 0;
|
||||
for(int i=0; i<in; i++) {
|
||||
dnnType *input = (dnnType*)reinterpret_cast<const dnnType*>(inputs[i]);
|
||||
int in_dim = c_in[i]*h*w;
|
||||
int part_in_dim = in_dim / this->groups;
|
||||
checkCuda( cudaMemcpyAsync(dstData + b*c*w*h + offset, input + b*c*w*h*groups + this->group_id*part_in_dim, part_in_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream) );
|
||||
offset += part_in_dim;
|
||||
}
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
return 0;
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return (6+MAX_INPUTS)*sizeof(int);
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, groups);
|
||||
tk::dnn::writeBUF(buf, group_id);
|
||||
tk::dnn::writeBUF(buf, in);
|
||||
for(int i=0; i<MAX_INPUTS; i++)
|
||||
tk::dnn::writeBUF(buf, c_in[i]);
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
static const int MAX_INPUTS = 4;
|
||||
int in;
|
||||
int c_in[MAX_INPUTS];
|
||||
int c, h, w;
|
||||
int groups, group_id;
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
static const int MAX_INPUTS = 4;
|
||||
int in;
|
||||
int c_in[MAX_INPUTS];
|
||||
int c, h, w;
|
||||
int groups, group_id;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class RouteRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
RouteRTPluginCreator() ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(RouteRTPluginCreator);
|
||||
};
|
||||
|
||||
@@ -1,75 +1,109 @@
|
||||
#ifndef _SHORTCUTRT_PLUGIN_H
|
||||
#define _SHORTCUTRT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
|
||||
class ShortcutRT : public IPlugin {
|
||||
|
||||
public:
|
||||
ShortcutRT(tk::dnn::dataDim_t bdim) {
|
||||
this->bc = bdim.c;
|
||||
this->bh = bdim.h;
|
||||
this->bw = bdim.w;
|
||||
}
|
||||
|
||||
~ShortcutRT(){
|
||||
|
||||
}
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW{inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]};
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
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 {
|
||||
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *srcDataBack = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
for(int b=0; b < batchSize; ++b)
|
||||
shortcutForward(srcDataBack + b*bc*bh*bw, dstData + b*c*h*w, 1, c, h, w, 1, 1, bc, bh, bw, 1, stream);
|
||||
|
||||
return 0;
|
||||
}
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <tkdnn.h>
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
namespace nvinfer1 {
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, bc);
|
||||
tk::dnn::writeBUF(buf, bh);
|
||||
tk::dnn::writeBUF(buf, bw);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
|
||||
}
|
||||
class ShortcutRT : public IPluginV2Ext {
|
||||
|
||||
public:
|
||||
ShortcutRT(int bc,int bh,int bw,int c,int h,int w ,bool mul);
|
||||
|
||||
~ShortcutRT();
|
||||
|
||||
ShortcutRT(const void *data, size_t length);
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims, int32_t nbOutputs,
|
||||
DataType const *inputTypes, DataType const *outputTypes, bool const *inputIsBroadcast,
|
||||
bool const *outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch (int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch (int32_t inputIndex) const NOEXCEPT override;
|
||||
|
||||
void attachToContext (cudnnContext *, cublasContext *, IGpuAllocator *) NOEXCEPT override;
|
||||
|
||||
void detachFromContext () NOEXCEPT override;
|
||||
|
||||
DataType getOutputDataType(int32_t index, nvinfer1::DataType const *inputTypes, int32_t nbInputs) const NOEXCEPT override;
|
||||
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override;
|
||||
|
||||
int c, h, w;
|
||||
int bc, bh, bw,bl;
|
||||
bool mul;
|
||||
tk::dnn::dataDim_t bDim;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
|
||||
class ShortcutRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ShortcutRTPluginCreator();
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
public:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ShortcutRTPluginCreator);
|
||||
|
||||
int c, h, w;
|
||||
int bc, bh, bw;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -1,66 +1,103 @@
|
||||
#ifndef _UPSAMPLERT_PLUGIN_H
|
||||
#define _UPSAMPLERT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
class UpsampleRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
UpsampleRT(int stride) {
|
||||
this->stride = stride;
|
||||
}
|
||||
class UpsampleRT : public IPluginV2Ext {
|
||||
|
||||
~UpsampleRT(){
|
||||
public:
|
||||
UpsampleRT(int stride,int c,int h,int w);
|
||||
|
||||
}
|
||||
UpsampleRT(const void *data, size_t length);
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
~UpsampleRT();
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return DimsCHW(inputs[0].d[0], inputs[0].d[1]*stride, inputs[0].d[2]*stride);
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
int initialize() override {
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
return 0;
|
||||
}
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
virtual void terminate() override {
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT 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 {
|
||||
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
|
||||
fill(dstData, batchSize*c*h*w*stride*stride, 0.0, stream);
|
||||
upsampleForward(srcData, dstData, batchSize, c, h, w, stride, 1, 1, stream);
|
||||
return 0;
|
||||
}
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, stride);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
int c, h, w, stride;
|
||||
};
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override ;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch (int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch (int32_t inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims, int32_t nbOutputs,
|
||||
DataType const *inputTypes, DataType const *outputTypes, bool const *inputIsBroadcast,
|
||||
bool const *outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void attachToContext (cudnnContext *, cublasContext *, IGpuAllocator *) NOEXCEPT override;
|
||||
|
||||
void detachFromContext () NOEXCEPT override;
|
||||
|
||||
DataType getOutputDataType (int32_t index, nvinfer1::DataType const *inputTypes, int32_t nbInputs) const NOEXCEPT override;
|
||||
|
||||
|
||||
int c, h, w, stride;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class UpsampleRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
UpsampleRTPluginCreator();
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(UpsampleRTPluginCreator);
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -1,143 +1,124 @@
|
||||
#ifndef _YOLORT_PLUGIN_H
|
||||
#define _YOLORT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include <vector>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <tkdnn.h>
|
||||
|
||||
#define YOLORT_CLASSNAME_W 256
|
||||
|
||||
class YoloRT : public IPlugin {
|
||||
namespace nvinfer1 {
|
||||
class YoloRT : public IPluginV2Ext {
|
||||
|
||||
public:
|
||||
YoloRT(int classes, int num,int c,int h,int w, int n_masks = 3, float scale_xy = 1,
|
||||
float nms_thresh = 0.45, int nms_kind = 0, int new_coords = 0);
|
||||
|
||||
YoloRT(const void *data, size_t length);
|
||||
|
||||
~YoloRT();
|
||||
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
public:
|
||||
YoloRT(int classes, int num, tk::dnn::Yolo *yolo = nullptr, int n_masks=3, float scale_xy=1, float nms_thresh=0.45, int nms_kind=0, int new_coords=0) {
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
this->classes = classes;
|
||||
this->num = num;
|
||||
this->n_masks = n_masks;
|
||||
this->scaleXY = scale_xy;
|
||||
this->nms_thresh = nms_thresh;
|
||||
this->nms_kind = nms_kind;
|
||||
this->new_coords = new_coords;
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
mask = new dnnType[n_masks];
|
||||
bias = new dnnType[num*n_masks*2];
|
||||
if(yolo != nullptr) {
|
||||
memcpy(mask, yolo->mask_h, sizeof(dnnType)*n_masks);
|
||||
memcpy(bias, yolo->bias_h, sizeof(dnnType)*num*n_masks*2);
|
||||
classesNames = yolo->classesNames;
|
||||
}
|
||||
}
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
~YoloRT(){
|
||||
|
||||
}
|
||||
|
||||
int getNbOutputs() const override {
|
||||
return 1;
|
||||
}
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
|
||||
return inputs[0];
|
||||
}
|
||||
|
||||
void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
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 {
|
||||
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
|
||||
for (int b = 0; b < batchSize; ++b){
|
||||
for(int n = 0; n < n_masks; ++n){
|
||||
int index = entry_index(b, n*w*h, 0);
|
||||
if (new_coords == 1){
|
||||
if (this->scaleXY != 1) scalAdd(dstData + index, 2 * w*h, this->scaleXY, -0.5*(this->scaleXY - 1), 1);
|
||||
}
|
||||
else{
|
||||
activationLOGISTICForward(srcData + index, dstData + index, 2*w*h, stream); //x,y
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
#elif NV_TENSORRT_MAJOR == 7
|
||||
int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
|
||||
#endif
|
||||
|
||||
if (this->scaleXY != 1) scalAdd(dstData + index, 2 * w*h, this->scaleXY, -0.5*(this->scaleXY - 1), 1);
|
||||
|
||||
index = entry_index(b, n*w*h, 4);
|
||||
activationLOGISTICForward(srcData + index, dstData + index, (1+classes)*w*h, stream);
|
||||
}
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2Ext *clone() const NOEXCEPT override;
|
||||
|
||||
DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
|
||||
|
||||
bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
|
||||
|
||||
bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
|
||||
|
||||
void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
|
||||
int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
|
||||
bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
|
||||
int32_t maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void detachFromContext() NOEXCEPT override;
|
||||
|
||||
|
||||
int c, h, w;
|
||||
int classes, num, n_masks;
|
||||
float scaleXY;
|
||||
float nms_thresh;
|
||||
int nms_kind;
|
||||
int new_coords;
|
||||
int NUM = 0;
|
||||
std::vector<std::string> classesNames;
|
||||
|
||||
|
||||
int entry_index(int batch, int location, int entry) {
|
||||
int n = location / (w * h);
|
||||
int loc = location % (w * h);
|
||||
return batch * c * h * w + n * w * h * (4 + classes + 1) + entry * w * h + loc;
|
||||
}
|
||||
|
||||
//std::cout<<"YOLO END\n";
|
||||
return 0;
|
||||
}
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
|
||||
};
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 8*sizeof(int) + 2*sizeof(float)+ n_masks*sizeof(dnnType) + num*n_masks*2*sizeof(dnnType) + YOLORT_CLASSNAME_W*classes*sizeof(char);
|
||||
}
|
||||
class YoloRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
YoloRTPluginCreator();
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
tk::dnn::writeBUF(buf, classes); //std::cout << "Classes :" << classes << std::endl;
|
||||
tk::dnn::writeBUF(buf, num); //std::cout << "Num : " << num << std::endl;
|
||||
tk::dnn::writeBUF(buf, n_masks); //std::cout << "N_Masks" << n_masks << std::endl;
|
||||
tk::dnn::writeBUF(buf, scaleXY); //std::cout << "ScaleXY :" << scaleXY << std::endl;
|
||||
tk::dnn::writeBUF(buf, nms_thresh); //std::cout << "nms_thresh :" << nms_thresh << std::endl;
|
||||
tk::dnn::writeBUF(buf, nms_kind); //std::cout << "nms_kind : " << nms_kind << std::endl;
|
||||
tk::dnn::writeBUF(buf, new_coords); //std::cout << "new_coords : " << new_coords << std::endl;
|
||||
tk::dnn::writeBUF(buf, c); //std::cout << "C : " << c << std::endl;
|
||||
tk::dnn::writeBUF(buf, h); //std::cout << "H : " << h << std::endl;
|
||||
tk::dnn::writeBUF(buf, w); //std::cout << "C : " << c << std::endl;
|
||||
for (int i = 0; i < n_masks; i++)
|
||||
{
|
||||
tk::dnn::writeBUF(buf, mask[i]); //std::cout << "mask[i] : " << mask[i] << std::endl;
|
||||
}
|
||||
for (int i = 0; i < n_masks * 2 * num; i++)
|
||||
{
|
||||
tk::dnn::writeBUF(buf, bias[i]); //std::cout << "bias[i] : " << bias[i] << std::endl;
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
// save classes names
|
||||
for(int i=0; i<classes; i++) {
|
||||
char tmp[YOLORT_CLASSNAME_W];
|
||||
strcpy(tmp, classesNames[i].c_str());
|
||||
for(int j=0; j<YOLORT_CLASSNAME_W; j++) {
|
||||
tk::dnn::writeBUF(buf, tmp[j]);
|
||||
}
|
||||
}
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
int c, h, w;
|
||||
int classes, num, n_masks;
|
||||
float scaleXY;
|
||||
float nms_thresh;
|
||||
int nms_kind;
|
||||
int new_coords;
|
||||
std::vector<std::string> classesNames;
|
||||
IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
dnnType *mask;
|
||||
dnnType *bias;
|
||||
IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
int entry_index(int batch, int location, int entry) {
|
||||
int n = location / (w*h);
|
||||
int loc = location % (w*h);
|
||||
return batch*c*h*w + n*w*h*(4+classes+1) + entry*w*h + loc;
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(YoloRTPluginCreator);
|
||||
};
|
||||
#endif
|
||||
+43
-1
@@ -6,11 +6,15 @@
|
||||
#include <fstream>
|
||||
#include <iomanip>
|
||||
#include <stdlib.h>
|
||||
#include <yaml-cpp/yaml.h>
|
||||
|
||||
|
||||
#include "cuda.h"
|
||||
#include "cuda_runtime_api.h"
|
||||
#include <cublas_v2.h>
|
||||
#include <cudnn.h>
|
||||
#include <NvInferVersion.h>
|
||||
|
||||
|
||||
#ifdef __linux__
|
||||
#include <unistd.h>
|
||||
@@ -20,8 +24,30 @@
|
||||
#include <chrono>
|
||||
|
||||
|
||||
|
||||
|
||||
#if NV_TENSORRT_MAJOR > 7
|
||||
#define NOEXCEPT noexcept
|
||||
#else
|
||||
#define NOEXCEPT
|
||||
#endif
|
||||
|
||||
|
||||
#define dnnType float
|
||||
|
||||
template<typename T> void writeBUF(char*& buffer, const T& val)
|
||||
{
|
||||
*reinterpret_cast<T*>(buffer) = val;
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
template<typename T> T readBUF(const char*& buffer)
|
||||
{
|
||||
T val = *reinterpret_cast<const T*>(buffer);
|
||||
buffer += sizeof(T);
|
||||
return val;
|
||||
}
|
||||
|
||||
|
||||
// Colored output
|
||||
#define COL_END "\033[0m"
|
||||
@@ -118,7 +144,7 @@ void printCenteredTitle(const char *title, char fill, int dim = 30);
|
||||
bool fileExist(const char *fname);
|
||||
void downloadWeightsifDoNotExist(const std::string& input_bin, const std::string& test_folder, const std::string& weights_url);
|
||||
void readBinaryFile(std::string fname, int size, dnnType** data_h, dnnType** data_d, int seek = 0);
|
||||
int checkResult(int size, dnnType *data_d, dnnType *correct_d, bool device = true, int limit = 10);
|
||||
int checkResult(int size, dnnType *data_d, dnnType *correct_d, bool device = true, int limit = 10, bool verbose=true);
|
||||
void printDeviceVector(int size, dnnType* vec_d, bool device = true);
|
||||
float getColor(const int c, const int x, const int max);
|
||||
void resize(int size, dnnType **data);
|
||||
@@ -135,4 +161,20 @@ static inline bool isCudaPointer(void *data) {
|
||||
cudaPointerAttributes attr;
|
||||
return cudaPointerGetAttributes(&attr, data) == 0;
|
||||
}
|
||||
|
||||
inline YAML::Node YAMLloadConf(const std::string& conf_file) {
|
||||
std::cerr<<"Loading YAML: "<<conf_file<<"\n";
|
||||
return YAML::LoadFile(conf_file);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
inline T YAMLgetConf(YAML::Node conf, std::string key, T defaultVal) {
|
||||
T val = defaultVal;
|
||||
if(conf && conf[key]) {
|
||||
val = conf[key].as<T>();
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
|
||||
#endif //UTILS_H
|
||||
|
||||
Reference in New Issue
Block a user