mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

YoloRT save bias, mask and clasesName into RT file

Browse Source
This commit is contained in:
Francesco Gatti
2022-03-30 20:46:51 +02:00
parent fa9db167b8
commit 5e71b99265
15 changed files with 104 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CenternetDetection.cpp
+1 -1
View File
@@ -3,7 +3,7 @@
namespace tk { namespace dnn {
bool CenternetDetection::init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh){
bool CenternetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
std::cout<<(tensor_path).c_str()<<"\n";
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() );
classes = n_classes;
src/MobilenetDetection.cpp
+1 -1
View File
@@ -126,7 +126,7 @@ float MobilenetDetection::iou(const tk::dnn::box &a, const tk::dnn::box &b){
return iou;
}
bool MobilenetDetection::init(const std::string& tensor_path, const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh){
bool MobilenetDetection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh){
std::cout<<(tensor_path).c_str()<<"\n";
netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str());
imageSize = netRT->input_dim.h;
src/NetworkRT.cpp
+16
View File
@@ -15,6 +15,9 @@
using namespace nvinfer1;
extern std::mutex gYoloPlugins_mutex;
extern std::vector<YoloRT*> gYoloPlugins;
// Logger for info/warning/errors
class Logger : public ILogger {
void log(Severity severity, const char* msg) NOEXCEPT override {
@@ -826,6 +829,12 @@ IPluginV2Layer* NetworkRT::convert_layer(ITensor *input, Yolo *l) {
mPluginAttributes.emplace_back(PluginField("nms_thresh",&l->nms_thresh,PluginFieldType::kFLOAT32,1));
mPluginAttributes.emplace_back(PluginField("nms_kins",&l->nsm_kind,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("new_coords",&l->new_coords,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("mask",l->mask_h,PluginFieldType::kFLOAT32,l->n_masks));
mPluginAttributes.emplace_back(PluginField("bias",l->bias_h,PluginFieldType::kFLOAT32,l->n_masks*2*l->num));
for(int i=0; i<l->classes; i++) {
mPluginAttributes.emplace_back(PluginField("class_name",l->classesNames[i].data(),PluginFieldType::kCHAR,l->classesNames[i].size()));
}
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC);
@@ -1001,7 +1010,14 @@ bool NetworkRT::deserialize(const char *filename) {
}
runtimeRT = createInferRuntime(loggerRT);
gYoloPlugins_mutex.lock();
gYoloPlugins.clear();
engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size);
yolo_plugins = gYoloPlugins;
gYoloPlugins.clear();
gYoloPlugins_mutex.unlock();
std::cout<<size<<std::endl;
//if (gieModelStream) delete [] gieModelStream;
src/Yolo3Detection.cpp
+17 -33
View File
@@ -3,7 +3,7 @@
namespace tk { namespace dnn {
bool Yolo3Detection::init(const std::string& tensor_path,const std::string& cfg_path,const std::string& name_path,const int n_classes, const int n_batches, const float conf_thresh) {
bool Yolo3Detection::init(const std::string& tensor_path, const int n_classes, const int n_batches, const float conf_thresh) {
//convert network to tensorRT
std::cout<<(tensor_path).c_str()<<"\n";
@@ -14,43 +14,27 @@ namespace tk { namespace dnn {
tk::dnn::dataDim_t idim = netRT->input_dim;
idim.n = nBatches;
std::vector<int> yolosLine = noYolosLine(cfg_path);
noYolos = yolosLine;
int channels,height,width;
loadYoloInitInfo(channels,width,height,cfg_path);
if(yolosLine.size() < 2 ) {
if(netRT->yolo_plugins.size() < 2 ) {
FatalError("this is not yolo3");
}
for(int i=0; i<netRT->yolo_plugins.size(); i++) {
nvinfer1::YoloRT *yRT = netRT->yolo_plugins[i];
classes = yRT->classes;
num = yRT->num;
nMasks = yRT->n_masks;
for(int i=0; i<noYolos.size(); i++) {
std::vector<float> maskTemp,anchorsTemp;
std::vector<std::string> classNamesTemp;
int nms_kind,coords,numTemp;
float nmsthresh;
loadYoloInfo(cfg_path,yolosLine[i],maskTemp,anchorsTemp,numTemp,classes,nmsthresh,nms_kind,coords);
classNamesTemp = darknetReadNames(name_path);
num = numTemp/maskTemp.size();
nMasks = maskTemp.size();
dnnType* maskTempF;
dnnType* biasTempF;
maskTempF = maskTemp.data();
biasTempF = anchorsTemp.data();
// make a yolo layer to interpret predictions
yolo[i] = new tk::dnn::Yolo(nullptr, classes, nMasks, ""); // yolo without input and bias
yolo[i]->mask_h = new dnnType[nMasks];
yolo[i]->bias_h = new dnnType[num*nMasks*2];
memcpy(yolo[i]->mask_h, maskTempF, sizeof(dnnType)*nMasks);
memcpy(yolo[i]->bias_h, biasTempF, sizeof(dnnType)*num*nMasks*2);
auto dim = netRT->engineRT->getBindingDimensions(i+1);
yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, dim.d[0], dim.d[1], dim.d[2]);
yolo[i]->classesNames = classNamesTemp;
yolo[i]->nms_thresh = nmsthresh;
yolo[i]->nsm_kind = (tk::dnn::Yolo::nmsKind_t) nms_kind;
yolo[i]->new_coords = coords;
memcpy(yolo[i]->mask_h, yRT->mask.data(), sizeof(dnnType)*nMasks);
memcpy(yolo[i]->bias_h, yRT->bias.data(), sizeof(dnnType)*num*nMasks*2);
yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w);
yolo[i]->classesNames = yRT->classesNames;
yolo[i]->nms_thresh = yRT->nms_thresh;
yolo[i]->nsm_kind = (tk::dnn::Yolo::nmsKind_t) yRT->nms_kind;
yolo[i]->new_coords = yRT->new_coords;
}
dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes);
@@ -112,12 +96,12 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
//get yolo outputs
if(noYolos.size() < 2){
if(netRT->yolo_plugins.size() < 2){
FatalError("YOLOS WRONG!!");
}
std::vector<float *> rt_out;
//dnnType *rt_out[netRT->pluginFactory->n_yolos];
for(int i=0; i<noYolos.size(); i++)
for(int i=0; i<netRT->yolo_plugins.size(); i++)
rt_out.push_back((dnnType*)netRT->buffersRT[i+1] + netRT->buffersDIM[i+1].tot()*bi);
float x_ratio = float(originalSize[bi].width) / float(netRT->input_dim.w);
@@ -125,7 +109,7 @@ void Yolo3Detection::postprocess(const int bi, const bool mAP){
// compute dets
nDets = 0;
for(int i=0; i<noYolos.size(); i++) {
for(int i=0; i<netRT->yolo_plugins.size(); i++) {
yolo[i]->dstData = rt_out[i];
yolo[i]->computeDetections(dets, nDets, netRT->input_dim.w, netRT->input_dim.h, confThreshold, yolo[i]->new_coords);
}
src/pluginsRT/YoloRT.cpp
+54 -1
View File
@@ -1,8 +1,13 @@
#include <tkDNN/pluginsRT/YoloRT.h>
#include <utility>
#include <mutex>
using namespace nvinfer1;
// used to retrive Yolo plugin during network deserialization
std::mutex gYoloPlugins_mutex;
std::vector<YoloRT*> gYoloPlugins;
std::vector<PluginField> YoloRTPluginCreator::mPluginAttributes;
PluginFieldCollection YoloRTPluginCreator::mFC{};
@@ -22,6 +27,10 @@ YoloRT::YoloRT(int classes, int num, int c,int h,int w,int n_masks, float scale_
this->nms_thresh = nms_thresh;
this->nms_kind = nms_kind;
this->new_coords = new_coords;
bias.clear();
mask.clear();
classesNames.clear();
}
YoloRT::YoloRT(const void *data, size_t length) {
@@ -36,7 +45,24 @@ YoloRT::YoloRT(const void *data, size_t length) {
c = readBUF<int>(buf);
h = readBUF<int>(buf);
w = readBUF<int>(buf);
mask.resize(n_masks);
for(int i=0; i<n_masks; i++)
mask[i] = readBUF<dnnType>(buf);
bias.resize(n_masks*2*num);
for(int i=0; i<n_masks*2*num; i++)
bias[i] = readBUF<dnnType>(buf);
// save classes names
classesNames.resize(classes);
for(int i=0; i<classes; i++) {
char tmp[YOLORT_CLASSNAME_W];
for(int j=0; j<YOLORT_CLASSNAME_W; j++)
tmp[j] = readBUF<char>(buf);
classesNames[i] = std::string(tmp);
}
assert(buf == bufCheck + length);
gYoloPlugins.push_back(this);
}
YoloRT::~YoloRT() {}
@@ -126,7 +152,7 @@ int32_t YoloRT::enqueue(int32_t batchSize, const void *const *inputs, void **out
size_t YoloRT::getSerializationSize() const NOEXCEPT {
return 8 * sizeof(int) + 2 * sizeof(float) ;
return 8 * sizeof(int) + 2 * sizeof(float) + n_masks*sizeof(dnnType) + num*n_masks*2*sizeof(dnnType) + YOLORT_CLASSNAME_W*classes*sizeof(char);
}
bool YoloRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
@@ -145,6 +171,19 @@ void YoloRT::serialize(void *buffer) const NOEXCEPT {
writeBUF(buf, c); //std::cout << "C : " << c << std::endl;
writeBUF(buf, h); //std::cout << "H : " << h << std::endl;
writeBUF(buf, w); //std::cout << "C : " << c << std::endl;
for (int i = 0; i < n_masks; i++)
writeBUF(buf, mask[i]); //std::cout << "mask[i] : " << mask[i] << std::endl;
for (int i = 0; i < n_masks * 2 * num; i++)
writeBUF(buf, bias[i]); //std::cout << "bias[i] : " << bias[i] << std::endl;
// 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++) {
writeBUF(buf, tmp[j]);
}
}
assert(buf == a + getSerializationSize());
}
@@ -171,6 +210,9 @@ void YoloRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
IPluginV2Ext *YoloRT::clone() const NOEXCEPT {
auto *p = new YoloRT(classes, num,c,h,w,n_masks, scaleXY, nms_thresh, nms_kind, new_coords);
p->mask = mask;
p->bias = bias;
p->classesNames = classesNames;
p->setPluginNamespace(mPluginNamespace.c_str());
return p;
}
@@ -235,6 +277,17 @@ IPluginV2Ext *YoloRTPluginCreator::createPlugin(const char *name, const PluginFi
int nms_kind = *(static_cast<const int*>(fields[8].data));
int new_coords = *(static_cast<const int*>(fields[9].data));
auto *pluginObj = new YoloRT(classes,num,c,h,w,n_masks,scaleXY,nmsThresh,nms_kind,new_coords);
// fill additional data
pluginObj->mask.resize(fields[10].length*sizeof(float));
memcpy(pluginObj->mask.data(), fields[10].data, fields[10].length*sizeof(float));
pluginObj->bias.resize(fields[11].length*sizeof(float));
memcpy(pluginObj->bias.data(), fields[11].data, fields[11].length*sizeof(float));
pluginObj->classesNames.resize(classes);
for(int i=0; i<classes; i++) {
pluginObj->classesNames[i].resize(fields[12+i].length);
memcpy(&pluginObj->classesNames[i][0], fields[12+i].data, fields[12+i].length*sizeof(char));
}
return pluginObj;
}