Francesco Gatti
319 lines
11 KiB
C++
319 lines
11 KiB
C++
#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{};
|
|
|
|
static const char* YOLORT_PLUGIN_VERSION{"1"};
|
|
static const char* YOLORT_PLUGIN_NAME{"YoloRT_tkDNN"};
|
|
|
|
YoloRT::YoloRT(int classes, int num, int c,int h,int w,int n_masks, float scale_xy,
|
|
float nms_thresh, int nms_kind,
|
|
int new_coords) {
|
|
this->c = c;
|
|
this->h = h;
|
|
this->w = w;
|
|
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;
|
|
|
|
bias.clear();
|
|
mask.clear();
|
|
classesNames.clear();
|
|
}
|
|
|
|
YoloRT::YoloRT(const void *data, size_t length) {
|
|
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
|
classes = readBUF<int>(buf);
|
|
num = readBUF<int>(buf);
|
|
n_masks = readBUF<int>(buf);
|
|
scaleXY = readBUF<float>(buf);
|
|
nms_thresh = readBUF<float>(buf);
|
|
nms_kind = readBUF<int>(buf);
|
|
new_coords = readBUF<int>(buf);
|
|
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() {}
|
|
|
|
int YoloRT::getNbOutputs() const NOEXCEPT {
|
|
return 1;
|
|
}
|
|
|
|
Dims YoloRT::getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT {
|
|
return inputs[0];
|
|
}
|
|
|
|
|
|
|
|
int YoloRT::initialize() NOEXCEPT {
|
|
return 0;
|
|
}
|
|
|
|
void YoloRT::terminate() NOEXCEPT {}
|
|
|
|
size_t YoloRT::getWorkspaceSize(int maxBatchSize) const NOEXCEPT {
|
|
return 0;
|
|
}
|
|
|
|
#if NV_TENSORRT_MAJOR > 7
|
|
int YoloRT::enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
|
cudaStream_t stream) NOEXCEPT {
|
|
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 < 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 (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);
|
|
}
|
|
}
|
|
}
|
|
|
|
//std::cout<<"YOLO END\n";
|
|
return 0;
|
|
}
|
|
#elif NV_TENSORRT_MAJOR == 7
|
|
int32_t YoloRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) {
|
|
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 < 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 (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);
|
|
}
|
|
}
|
|
}
|
|
|
|
//std::cout<<"YOLO END\n";
|
|
return 0;
|
|
}
|
|
#endif
|
|
|
|
|
|
size_t YoloRT::getSerializationSize() const NOEXCEPT {
|
|
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 {
|
|
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
|
}
|
|
|
|
void YoloRT::serialize(void *buffer) const NOEXCEPT {
|
|
char *buf = reinterpret_cast<char *>(buffer), *a = buf;
|
|
writeBUF(buf, classes); //std::cout << "Classes :" << classes << std::endl;
|
|
writeBUF(buf, num); //std::cout << "Num : " << num << std::endl;
|
|
writeBUF(buf, n_masks); //std::cout << "N_Masks" << n_masks << std::endl;
|
|
writeBUF(buf, scaleXY); //std::cout << "ScaleXY :" << scaleXY << std::endl;
|
|
writeBUF(buf, nms_thresh); //std::cout << "nms_thresh :" << nms_thresh << std::endl;
|
|
writeBUF(buf, nms_kind); //std::cout << "nms_kind : " << nms_kind << std::endl;
|
|
writeBUF(buf, new_coords); //std::cout << "new_coords : " << new_coords << std::endl;
|
|
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());
|
|
}
|
|
|
|
const char *YoloRT::getPluginType() const NOEXCEPT {
|
|
return YOLORT_PLUGIN_NAME;
|
|
}
|
|
|
|
const char *YoloRT::getPluginVersion() const NOEXCEPT {
|
|
return YOLORT_PLUGIN_VERSION;
|
|
}
|
|
|
|
void YoloRT::destroy() NOEXCEPT {
|
|
delete this;
|
|
}
|
|
|
|
const char *YoloRT::getPluginNamespace() const NOEXCEPT {
|
|
return mPluginNamespace.c_str();
|
|
}
|
|
|
|
void YoloRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
|
|
mPluginNamespace = pluginNamespace;
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
DataType YoloRT::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const NOEXCEPT {
|
|
return DataType::kFLOAT;
|
|
}
|
|
|
|
void YoloRT::attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
|
|
IGpuAllocator *gpuAllocator) NOEXCEPT {
|
|
|
|
}
|
|
|
|
void YoloRT::configurePlugin(const Dims *inputDims, int32_t nbInputs, const Dims *outputDims, int32_t nbOutputs,
|
|
const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
|
|
const bool *outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT {
|
|
|
|
}
|
|
|
|
bool YoloRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const NOEXCEPT {
|
|
return false;
|
|
}
|
|
|
|
bool YoloRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT {
|
|
return false;
|
|
}
|
|
|
|
void YoloRT::detachFromContext() NOEXCEPT {
|
|
|
|
}
|
|
|
|
YoloRTPluginCreator::YoloRTPluginCreator() {
|
|
mPluginAttributes.clear();
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
void YoloRTPluginCreator::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
|
|
mPluginNamespace = pluginNamespace;
|
|
}
|
|
|
|
const char *YoloRTPluginCreator::getPluginNamespace() const NOEXCEPT {
|
|
return mPluginNamespace.c_str();
|
|
}
|
|
|
|
IPluginV2Ext *YoloRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT {
|
|
auto *pluginObj = new YoloRT(serialData,serialLength);
|
|
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
|
return pluginObj;
|
|
}
|
|
|
|
IPluginV2Ext *YoloRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
|
|
const PluginField *fields = fc->fields;
|
|
int classes = *(static_cast<const int *>(fields[0].data));
|
|
int num = *(static_cast<const int *>(fields[1].data));
|
|
int c = *(static_cast<const int *>(fields[2].data));
|
|
int h = *(static_cast<const int *>(fields[3].data));
|
|
int w = *(static_cast<const int *>(fields[4].data));
|
|
int n_masks = *(static_cast<const int *>(fields[5].data));
|
|
dnnType scaleXY = *(static_cast<const float*>(fields[6].data));
|
|
dnnType nmsThresh = *(static_cast<const float*>(fields[7].data));
|
|
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;
|
|
}
|
|
|
|
const char *YoloRTPluginCreator::getPluginName() const NOEXCEPT {
|
|
return YOLORT_PLUGIN_NAME;
|
|
}
|
|
|
|
const char *YoloRTPluginCreator::getPluginVersion() const NOEXCEPT {
|
|
return YOLORT_PLUGIN_VERSION;
|
|
}
|
|
|
|
const PluginFieldCollection *YoloRTPluginCreator::getFieldNames() NOEXCEPT {
|
|
return &mFC;
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|