tkDNN can now deserialize tensorrt-8 engine (both through test_* and trtexec)
but demo has issues in yolo::computeDetections
This commit is contained in:
Harshvardhan Chandirasekar
@@ -1,148 +1,84 @@
|
||||
#include "NvInfer.h"
|
||||
#include "../kernels.h"
|
||||
#include <cassert>
|
||||
#include <vector>
|
||||
|
||||
class ActivationLeakyRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class ActivationLeakyRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationLeakyRT(float s) { slope = s; }
|
||||
public:
|
||||
explicit ActivationLeakyRT(float s);
|
||||
|
||||
ActivationLeakyRT(const void *data, size_t length)
|
||||
{
|
||||
std::cout<<"DESERIALIZE LEAKYRT"<<std::endl;
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
slope = readBUF<float>(buf);
|
||||
size = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
ActivationLeakyRT(const void *data, size_t length);
|
||||
|
||||
}
|
||||
~ActivationLeakyRT();
|
||||
|
||||
~ActivationLeakyRT() {}
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override { return 1; }
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override {
|
||||
return inputs[0];
|
||||
}
|
||||
void
|
||||
configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override
|
||||
{
|
||||
assert(type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
size = 1;
|
||||
for (int i = 0; i < outputDims[0].nbDims; i++)
|
||||
size *= outputDims[0].d[i];
|
||||
}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
int initialize() NOEXCEPT override { return 0; }
|
||||
void terminate() NOEXCEPT override {}
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override {
|
||||
return 0;
|
||||
}
|
||||
int enqueue(int batchSize, void const *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
|
||||
virtual int enqueue(int batchSize, void const *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override {
|
||||
activationLEAKYForward(
|
||||
(dnnType *) reinterpret_cast<const dnnType *>(inputs[0]),
|
||||
reinterpret_cast<dnnType *>(outputs[0]), batchSize * size, slope,
|
||||
stream);
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 1 * sizeof(int) + 1 * sizeof(float);
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
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());
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override {
|
||||
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override {
|
||||
return "ActivationLeakyRT_tkDNN";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override {
|
||||
return "1";
|
||||
}
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override;
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override {
|
||||
ActivationLeakyRT *p = new ActivationLeakyRT(slope);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
int size;
|
||||
float slope;
|
||||
|
||||
int size;
|
||||
float slope;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
class ActivationLeakyRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationLeakyRTPluginCreator();
|
||||
|
||||
class ActivationLeakyRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationLeakyRTPluginCreator() {
|
||||
mPluginAttributes.emplace_back(
|
||||
PluginField("slope", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override {
|
||||
ActivationLeakyRT *pluginObj = new ActivationLeakyRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override {
|
||||
const PluginField *fields = fc->fields;
|
||||
assert(fc->nbFields == 1);
|
||||
assert(fields[0].type == PluginFieldType::kFLOAT32);
|
||||
float slope = *(static_cast<const float *>(fields[0].data));
|
||||
ActivationLeakyRT *pluginObj = new ActivationLeakyRT(slope);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ActivationLeakyRT_tkDNN";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLeakyRTPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLeakyRTPluginCreator);
|
||||
};
|
||||
@@ -1,147 +1,83 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class ActivationLogisticRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
ActivationLogisticRT() {
|
||||
class ActivationLogisticRT : public IPluginV2 {
|
||||
|
||||
}
|
||||
public:
|
||||
ActivationLogisticRT() ;
|
||||
|
||||
ActivationLogisticRT(const void *data, size_t length)
|
||||
{
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
size = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
ActivationLogisticRT(const void *data, size_t length) ;
|
||||
|
||||
}
|
||||
~ActivationLogisticRT() ;
|
||||
|
||||
~ActivationLogisticRT(){
|
||||
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 configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT 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;
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
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(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
activationLOGISTICForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]), batchSize*size, stream);
|
||||
return 0;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 1*sizeof(int);
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT override {
|
||||
char *buf = reinterpret_cast<char*>(buffer);
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override {
|
||||
return "ActivationLogisticRT_tkDNN";
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override {
|
||||
return "1";
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
int size;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
class ActivationLogisticRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationLogisticRTPluginCreator() ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
//todo assert;
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override{
|
||||
ActivationLogisticRT *p = new ActivationLogisticRT();
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
int size;
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
class ActivationLogisticRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
ActivationLogisticRTPluginCreator(){
|
||||
mPluginAttributes.clear();
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override {
|
||||
ActivationLogisticRT *pluginObj = new ActivationLogisticRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override {
|
||||
ActivationLogisticRT *pluginObj = new ActivationLogisticRT();
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ActivationLogisticRT_tkDNN";
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLogisticRTPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationLogisticRTPluginCreator);
|
||||
};
|
||||
@@ -1,135 +1,79 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
class ActivationMishRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class ActivationMishRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ActivationMishRT() {}
|
||||
public:
|
||||
ActivationMishRT() ;
|
||||
|
||||
~ActivationMishRT() {}
|
||||
~ActivationMishRT() ;
|
||||
|
||||
ActivationMishRT(const void *data, size_t length) {
|
||||
std::cout<<"DESERIALIZE MISH"<<std::endl;
|
||||
const char *buf = reinterpret_cast<const char *>(data), *bufCheck = buf;
|
||||
size = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
ActivationMishRT(const void *data, size_t length) ;
|
||||
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override { return 1; }
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override { return inputs[0]; }
|
||||
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 {
|
||||
assert(format == PluginFormat::kLINEAR);
|
||||
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() NOEXCEPT override { return 0; }
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0; }
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override {
|
||||
activationMishForward((dnnType *) reinterpret_cast<const dnnType *>(inputs[0]),
|
||||
reinterpret_cast<dnnType *>(outputs[0]), batchSize * size, stream);
|
||||
return 0;
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 1 * sizeof(int);
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
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());
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override {
|
||||
return "ActivationMishRT_tkDNN";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override {
|
||||
return "1";
|
||||
}
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
void setPluginNamespace(const char *plguinNamespace) NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *plguinNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = plguinNamespace;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override {
|
||||
ActivationMishRT *p = new ActivationMishRT();
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
int size;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
int size;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
class ActivationMishRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationMishRTPluginCreator() ;
|
||||
|
||||
class ActivationMishRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationMishRTPluginCreator() {
|
||||
mPluginAttributes.clear();
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char* name,const void* serialData,size_t serialLength) NOEXCEPT override{
|
||||
ActivationMishRT *pluginObj = new ActivationMishRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
ActivationMishRT *pluginObj = new ActivationMishRT();
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ActivationMishRT_tkDNN";
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationMishRTPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationMishRTPluginCreator);
|
||||
};
|
||||
@@ -1,150 +1,78 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
namespace nvinfer1 {
|
||||
class ActivationReLUCeiling : public IPluginV2 {
|
||||
|
||||
class ActivationReLUCeiling : public IPluginV2 {
|
||||
public:
|
||||
explicit ActivationReLUCeiling(const float ceiling) ;
|
||||
|
||||
public:
|
||||
ActivationReLUCeiling(const float ceiling) {
|
||||
this->ceiling = ceiling;
|
||||
}
|
||||
~ActivationReLUCeiling() ;
|
||||
|
||||
~ActivationReLUCeiling() {
|
||||
ActivationReLUCeiling(const void *data, size_t length) ;
|
||||
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
ActivationReLUCeiling(const void *data, size_t length) {
|
||||
std::cout<<"RELU CEILING DESERIALIZE"<<std::endl;
|
||||
const char *buf = reinterpret_cast<const char *>(data), *bufCheck = buf;
|
||||
ceiling = readBUF<float>(buf);
|
||||
size = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
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 configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override {
|
||||
assert(type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
size = 1;
|
||||
for (int i = 0; i < outputDims[0].nbDims; i++)
|
||||
size *= outputDims[0].d[i];
|
||||
}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override { return 0; }
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override {
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override {
|
||||
activationReLUCeilingForward((dnnType *) reinterpret_cast<const dnnType *>(inputs[0]),
|
||||
reinterpret_cast<dnnType *>(outputs[0]), batchSize * size, ceiling, stream);
|
||||
return 0;
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 1 * sizeof(int) + 1 * sizeof(float);
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void *buffer) const NOEXCEPT override {
|
||||
char *buf = reinterpret_cast<char *>(buffer), *a = buf;
|
||||
tk::dnn::writeBUF(buf, ceiling);
|
||||
tk::dnn::writeBUF(buf, size);
|
||||
assert(buf = a + getSerializationSize());
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override {
|
||||
ActivationReLUCeiling *p = new ActivationReLUCeiling(ceiling);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override {
|
||||
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override { delete this; };
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
int size;
|
||||
float ceiling;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override {
|
||||
return "ActivationReLUCeilingRT_tkDNN";
|
||||
}
|
||||
class ActivationReLUCeilingPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationReLUCeilingPluginCreator() ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override {
|
||||
return "1";
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
int size;
|
||||
float ceiling;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
class ActivationReLUCeilingPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ActivationReLUCeilingPluginCreator() {
|
||||
mPluginAttributes.emplace_back(PluginField("ceiling", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
public:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override {
|
||||
ActivationReLUCeiling *pluginObj = new ActivationReLUCeiling(serialData, serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char *name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
float ceiling = *(static_cast<const float *>(fields[0].data));
|
||||
ActivationReLUCeiling *pluginObj = new ActivationReLUCeiling(ceiling);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ActivationReLUCeilingRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
public:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationReLUCeilingPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(ActivationReLUCeilingPluginCreator);
|
||||
};
|
||||
@@ -1,360 +1,115 @@
|
||||
#ifndef _DEFORMABLECONVRT_PLUGIN_H
|
||||
#define _DEFORMABLECONVRT_PLUGIN_H
|
||||
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <tkdnn.h>
|
||||
|
||||
class DeformableConvRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class DeformableConvRT : public IPluginV2 {
|
||||
|
||||
|
||||
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);
|
||||
|
||||
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;
|
||||
this->defRT = deformable;
|
||||
~DeformableConvRT();
|
||||
|
||||
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;
|
||||
|
||||
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) {
|
||||
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(const void *data, size_t length) ;
|
||||
|
||||
~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 NOEXCEPT override ;
|
||||
|
||||
DeformableConvRT(const void *data,size_t length){
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
chunk_dim = readBUF<int>(buf);
|
||||
kh = readBUF<int>(buf);
|
||||
kw = readBUF<int>(buf);
|
||||
sh = readBUF<int>(buf);
|
||||
sw = readBUF<int>(buf);
|
||||
ph = readBUF<int>(buf);
|
||||
pw = readBUF<int>(buf);
|
||||
deformableGroup = readBUF<int>(buf);
|
||||
i_n = readBUF<int>(buf);
|
||||
i_c = readBUF<int>(buf);
|
||||
i_h = readBUF<int>(buf);
|
||||
i_w = readBUF<int>(buf);
|
||||
o_n = readBUF<int>(buf);
|
||||
o_c = readBUF<int>(buf);
|
||||
o_h = readBUF<int>(buf);
|
||||
o_w = readBUF<int>(buf);
|
||||
dnnType *aus = new dnnType[chunk_dim*2];
|
||||
for(int i=0;i<chunk_dim*2;i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda(cudaMemcpy(offset,aus,sizeof(dnnType)*2*chunk_dim,cudaMemcpyHostToDevice));
|
||||
free(aus);
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
aus = new dnnType[chunk_dim];
|
||||
for(int i=0;i<chunk_dim;i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda(cudaMemcpy(mask,aus,sizeof(dnnType)*chunk_dim,cudaMemcpyHostToDevice));
|
||||
free(aus);
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
aus = new dnnType[i_c*o_c*kh*kw*1];
|
||||
for(int i=0;i<(i_c*o_c*kh*kw*1);i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda(cudaMemcpy(data_d,aus,sizeof(dnnType)*(i_c*o_c*kh*kw*1),cudaMemcpyHostToDevice));
|
||||
free(aus);
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
aus = new dnnType[o_c];
|
||||
for(int i=0; i < o_c; i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda( cudaMemcpy(bias2_d, aus, sizeof(dnnType)*o_c, cudaMemcpyHostToDevice) );
|
||||
free(aus);
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
aus = new dnnType[height_ones * width_ones];
|
||||
for(int i=0; i<height_ones * width_ones; i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda( cudaMemcpy(ones_d1, aus, sizeof(dnnType)*height_ones * width_ones, cudaMemcpyHostToDevice) );
|
||||
free(aus);
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
aus = new dnnType[dim_ones];
|
||||
for(int i=0; i<dim_ones; i++)
|
||||
aus[i] = readBUF<dnnType>(buf);
|
||||
checkCuda( cudaMemcpy(ones_d2, aus, sizeof(dnnType)*dim_ones, cudaMemcpyHostToDevice) );
|
||||
free(aus);
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {return 1;}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{defRT->output_dim.c, defRT->output_dim.h, defRT->output_dim.w};
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT override { }
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override {return 0;}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override { }
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
// 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;
|
||||
}
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT 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);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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());
|
||||
}
|
||||
|
||||
void destroy() NOEXCEPT override {delete this;}
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
//todo assert
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "DeformableConvRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override{
|
||||
DeformableConvRT *p = new DeformableConvRT(chunk_dim,kh,kw,sh,sw,ph,pw,deformableGroup,i_n,i_c,i_h,i_w,o_n,o_c,o_h,o_w,defRT);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
|
||||
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;
|
||||
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;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class DeformableConvRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
DeformableConvRTPluginCreator();
|
||||
|
||||
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;
|
||||
};
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class DeformableConvRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
DeformableConvRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("chunk_dim",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("kh",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("kw",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("sh",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("sw",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("ph",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("pw",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("deformableGroup",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("i_n",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("i_c",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("i_h",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("i_w",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_n",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_c",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_h",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_w",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("defRT",nullptr,PluginFieldType::kUNKNOWN,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
DeformableConvRT *pluginObj = new DeformableConvRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
int chunk_dim = *(static_cast<const int *>(fields[0].data));
|
||||
int kh = *(static_cast<const int *>(fields[1].data));
|
||||
int kw = *(static_cast<const int *>(fields[2].data));
|
||||
int sh = *(static_cast<const int *>(fields[3].data));
|
||||
int sw = *(static_cast<const int *>(fields[4].data));
|
||||
int ph = *(static_cast<const int *>(fields[5].data));
|
||||
int pw = *(static_cast<const int *>(fields[6].data));
|
||||
int deformableGroup = *(static_cast<const int *>(fields[7].data));
|
||||
int i_n = *(static_cast<const int *>(fields[8].data));
|
||||
int i_c = *(static_cast<const int *>(fields[9].data));
|
||||
int i_h = *(static_cast<const int *>(fields[10].data));
|
||||
int i_w = *(static_cast<const int *>(fields[11].data));
|
||||
int o_n = *(static_cast<const int *>(fields[12].data));
|
||||
int o_c = *(static_cast<const int *>(fields[13].data));
|
||||
int o_h = *(static_cast<const int *>(fields[14].data));
|
||||
int o_w = *(static_cast<const int *>(fields[14].data));
|
||||
DeformConv2d *defRT = const_cast<DeformConv2d *>(static_cast<const DeformConv2d *>(fields[15].data));
|
||||
DeformableConvRT *pluginObj = new DeformableConvRT(chunk_dim,kh,kw,sh,sw,ph,pw,deformableGroup,i_n,i_c,i_h,i_w,o_n,o_c,o_h,o_w,defRT);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "DeformableConvRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
|
||||
#endif
|
||||
|
||||
@@ -1,162 +1,81 @@
|
||||
#include<cassert>
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
namespace nvinfer1 {
|
||||
class FlattenConcatRT : public IPluginV2 {
|
||||
|
||||
class FlattenConcatRT : public IPluginV2 {
|
||||
public:
|
||||
FlattenConcatRT() ;
|
||||
|
||||
public:
|
||||
FlattenConcatRT() {
|
||||
stat = cublasCreate(&handle);
|
||||
if (stat != CUBLAS_STATUS_SUCCESS) {
|
||||
printf ("CUBLAS initialization failed\n");
|
||||
return;
|
||||
}
|
||||
}
|
||||
FlattenConcatRT(const void *data, size_t length) ;
|
||||
|
||||
FlattenConcatRT(const void *data,size_t length){
|
||||
const char *buf = reinterpret_cast<const char *>(data),*bufCheck=buf;
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
rows = readBUF<int>(buf);
|
||||
cols = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
~FlattenConcatRT() ;
|
||||
|
||||
~FlattenConcatRT(){
|
||||
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 Dims3{ inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2], 1, 1};
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT 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() NOEXCEPT override {return 0;}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override { checkERROR(cublasDestroy(handle));}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override {return 0;}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT 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 ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 5*sizeof(int);
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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 *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "FlattenConcatRT_tkDNN";
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
int c, h, w;
|
||||
int rows, cols;
|
||||
cublasStatus_t stat;
|
||||
cublasHandle_t handle;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
class FlattenConcatRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
FlattenConcatRTPluginCreator() ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override {
|
||||
FlattenConcatRT *p = new FlattenConcatRT();
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
int c, h, w;
|
||||
int rows, cols;
|
||||
cublasStatus_t stat;
|
||||
cublasHandle_t handle;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
class FlattenConcatRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
FlattenConcatRTPluginCreator(){
|
||||
mPluginAttributes.clear();
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
FlattenConcatRT *pluginObj = new FlattenConcatRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
IPluginV2 *createPlugin(const char *name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
FlattenConcatRT *pluginObj = new FlattenConcatRT();
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "FlattenConcatRT_tkDNN";
|
||||
}
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(FlattenConcatRTPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(FlattenConcatRTPluginCreator);
|
||||
};
|
||||
@@ -1,182 +1,89 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
|
||||
class MaxPoolFixedSizeRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class MaxPoolFixedSizeRT : public IPluginV2 {
|
||||
|
||||
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;
|
||||
}
|
||||
public:
|
||||
MaxPoolFixedSizeRT(int c, int h, int w, int n, int strideH, int strideW, int winSize, int padding) ;
|
||||
|
||||
MaxPoolFixedSizeRT(const void *data,size_t length){
|
||||
const char *buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
n = readBUF<int>(buf);
|
||||
stride_H = readBUF<int>(buf);
|
||||
stride_W = readBUF<int>(buf);
|
||||
winSize = readBUF<int>(buf);
|
||||
padding = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
MaxPoolFixedSizeRT(const void *data, size_t length) ;
|
||||
|
||||
~MaxPoolFixedSizeRT(){
|
||||
}
|
||||
~MaxPoolFixedSizeRT() ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{this->c, this->h, this->w};
|
||||
}
|
||||
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 {
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override {return 0;}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT 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;
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 8*sizeof(int);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
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());
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
//todo assert
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "MaxPoolingFixedSizeRT_tkDNN";
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override{
|
||||
MaxPoolFixedSizeRT *p = new MaxPoolFixedSizeRT(c,h,w,n,stride_H,stride_W,winSize,padding);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
|
||||
int n, c, h, w;
|
||||
int stride_H, stride_W;
|
||||
int winSize;
|
||||
int padding;
|
||||
int n, c, h, w;
|
||||
int stride_H, stride_W;
|
||||
int winSize;
|
||||
int padding;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class MaxPoolFixedSizeRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
MaxPoolFixedSizeRTPluginCreator() ;
|
||||
|
||||
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(MaxPoolFixedSizeRTPluginCreator);
|
||||
};
|
||||
|
||||
class MaxPoolFixedSizeRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
MaxPoolFixedSizeRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("c",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("h",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("w",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("n",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("stride_H",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("stride_W",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("winSize",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("padding",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
MaxPoolFixedSizeRT *pluginObj = new MaxPoolFixedSizeRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char *name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
//todo assert
|
||||
int c = *(static_cast<const int *>(fields[0].data));
|
||||
int h = *(static_cast<const int *>(fields[1].data));
|
||||
int w = *(static_cast<const int *>(fields[2].data));
|
||||
int n = *(static_cast<const int *>(fields[3].data));
|
||||
int stride_H = *(static_cast<const int *>(fields[4].data));
|
||||
int stride_W = *(static_cast<const int *>(fields[5].data));
|
||||
int winSize = *(static_cast<const int *>(fields[6].data));
|
||||
int padding = *(static_cast<const int *>(fields[7].data));
|
||||
MaxPoolFixedSizeRT *pluginObj = new MaxPoolFixedSizeRT(c,h,w,n,stride_H,stride_W,winSize,padding);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "MaxPoolingFixedSizeRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(MaxPoolFixedSizeRTPluginCreator);
|
||||
|
||||
@@ -1,188 +1,95 @@
|
||||
#ifndef _REGIONRT_PLUGIN_H
|
||||
#define _REGIONRT_PLUGIN_H
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class RegionRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class RegionRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
RegionRT(int classes, int coords, int num) {
|
||||
this->classes = classes;
|
||||
this->coords = coords;
|
||||
this->num = num;
|
||||
}
|
||||
public:
|
||||
RegionRT(int classes, int coords, int num);
|
||||
|
||||
~RegionRT(){
|
||||
~RegionRT() ;
|
||||
|
||||
}
|
||||
RegionRT(const void *data, size_t length) ;
|
||||
|
||||
RegionRT(const void *data,size_t length){
|
||||
const char *buf = reinterpret_cast<const char*>(data),*bufCheck=buf;
|
||||
classes = readBUF<int>(buf);
|
||||
coords = readBUF<int>(buf);
|
||||
num = readBUF<int>(buf);
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck+length);
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return inputs[0];
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
|
||||
int initialize() NOEXCEPT override { return 0; }
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual void terminate() NOEXCEPT override { }
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0; }
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT 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 enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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());
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "RegionRT_tkDNN";
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override {delete this;}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override{
|
||||
RegionRT *p = new RegionRT(classes,coords,num);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
int c, h, w;
|
||||
int classes, coords, num;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
};
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
class RegionRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
RegionRTPluginCreator();
|
||||
|
||||
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(RegionRTPluginCreator);
|
||||
};
|
||||
|
||||
class RegionRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
RegionRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("classes",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("coords",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("num",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char* name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
RegionRT *pluginObj = new RegionRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
assert(fc->nbFields == 3);
|
||||
assert(fields[0].type == PluginFieldType::kINT32);
|
||||
assert(fields[1].type == PluginFieldType::kINT32);
|
||||
assert(fields[2].type == PluginFieldType::kINT32);
|
||||
int classes = *(static_cast<const int*>(fields[0].data));
|
||||
int coords = *(static_cast<const int*>(fields[1].data));
|
||||
int num = *(static_cast<const int*>(fields[2].data));
|
||||
RegionRT *pluginObj = new RegionRT(classes,coords,num);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "RegionRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(RegionRTPluginCreator);
|
||||
#endif
|
||||
|
||||
|
||||
@@ -1,143 +1,83 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
class ReorgRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
class ReorgRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ReorgRT(int stride) {
|
||||
this->stride = stride;
|
||||
}
|
||||
public:
|
||||
ReorgRT(int stride);
|
||||
|
||||
~ReorgRT(){
|
||||
~ReorgRT();
|
||||
|
||||
}
|
||||
ReorgRT(const void *data, size_t length);
|
||||
|
||||
ReorgRT(const void* data,size_t length){
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
stride = readBUF<int>(buf);
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {return 1;}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{inputs[0].d[0]*stride*stride, inputs[0].d[1]/stride, inputs[0].d[2]/stride};
|
||||
}
|
||||
void
|
||||
configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
int initialize() NOEXCEPT override { return 0;}
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
|
||||
reorgForward((dnnType*)reinterpret_cast<const dnnType*>(inputs[0]),
|
||||
reinterpret_cast<dnnType*>(outputs[0]),
|
||||
batchSize, c, h, w, stride, stream);
|
||||
return 0;
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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());
|
||||
}
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{return true;}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "ReorgRT_tkDNN";
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char* getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
void destroy() NOEXCEPT override{ delete this;}
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override{
|
||||
ReorgRT *p = new ReorgRT(stride);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
int c, h, w, stride;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2 *clone() const 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;
|
||||
|
||||
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(ReorgRTPluginCreator);
|
||||
};
|
||||
|
||||
class ReorgRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
ReorgRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("stride",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2* deserializePlugin(const char* name,const void* serialData,size_t serialLength) NOEXCEPT override{
|
||||
ReorgRT *pluginObj = new ReorgRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection* fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
assert(fc->nbFields == 1);
|
||||
assert(fields[0].type == PluginFieldType::kINT32);
|
||||
int stride = *(static_cast<const int *>(fields[0].data));
|
||||
ReorgRT *pluginObj = new ReorgRT(stride);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ReorgRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ReorgRTPluginCreator);
|
||||
|
||||
|
||||
@@ -1,148 +1,85 @@
|
||||
#ifndef _RESHAPERT_PLUGIN_H
|
||||
#define _RESHAPERT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <tkdnn.h>
|
||||
using namespace tk::dnn;
|
||||
|
||||
class ReshapeRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
ReshapeRT(dataDim_t newDim) {
|
||||
new_dim = newDim;
|
||||
n = new_dim.n;
|
||||
c = new_dim.c;
|
||||
h = new_dim.h;
|
||||
w = new_dim.w;
|
||||
}
|
||||
namespace nvinfer1 {
|
||||
class ReshapeRT : public IPluginV2 {
|
||||
|
||||
ReshapeRT(const void *data,size_t length){
|
||||
const char *buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
new_dim.n = readBUF<int>(buf);
|
||||
new_dim.c = readBUF<int>(buf);
|
||||
new_dim.h = readBUF<int>(buf);
|
||||
new_dim.w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
public:
|
||||
explicit ReshapeRT(dataDim_t newDim) ;
|
||||
|
||||
~ReshapeRT(){
|
||||
ReshapeRT(const void *data, size_t length) ;
|
||||
|
||||
}
|
||||
~ReshapeRT() ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{ c,h,w};
|
||||
}
|
||||
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 {
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override {return 0;}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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());
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
//todo assert
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "ReshapeRT_tkDNN";
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
int n, c, h, w;
|
||||
dataDim_t new_dim;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override{
|
||||
ReshapeRT *p = new ReshapeRT(new_dim);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
class ReshapeRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ReshapeRTPluginCreator() ;
|
||||
|
||||
int n, c, h, w;
|
||||
dataDim_t new_dim;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
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(ReshapeRTPluginCreator);
|
||||
};
|
||||
|
||||
class ReshapeRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
ReshapeRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("new_dim",nullptr,PluginFieldType::kUNKNOWN,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char* name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
ReshapeRT *pluginObj = new ReshapeRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
dataDim_t newDim = *(static_cast<const dataDim_t *>(fields[0].data));
|
||||
ReshapeRT *pluginObj = new ReshapeRT(newDim);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ReshapeRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ReshapeRTPluginCreator);
|
||||
#endif
|
||||
@@ -1,165 +1,87 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <utils.h>
|
||||
|
||||
class ResizeLayerRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
ResizeLayerRT(int c, int h, int w) {
|
||||
o_c = c;
|
||||
o_h = h;
|
||||
o_w = w;
|
||||
}
|
||||
class ResizeLayerRT : public IPluginV2 {
|
||||
|
||||
ResizeLayerRT(const void *data,size_t length){
|
||||
const char *buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
o_c = readBUF<int>(buf);
|
||||
o_h = readBUF<int>(buf);
|
||||
o_w = readBUF<int>(buf);
|
||||
i_c = readBUF<int>(buf);
|
||||
i_h = readBUF<int>(buf);
|
||||
i_w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
public:
|
||||
ResizeLayerRT(int c, int h, int w) ;
|
||||
|
||||
~ResizeLayerRT(){
|
||||
}
|
||||
ResizeLayerRT(const void *data, size_t length) ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
~ResizeLayerRT() ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{o_c, o_h, o_w};
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT override {
|
||||
i_c = inputDims[0].d[0];
|
||||
i_h = inputDims[0].d[1];
|
||||
i_w = inputDims[0].d[2];
|
||||
}
|
||||
|
||||
int initialize() NOEXCEPT override {return 0;}
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT 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;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT override {
|
||||
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
tk::dnn::writeBUF(buf, o_c);
|
||||
tk::dnn::writeBUF(buf, o_h);
|
||||
tk::dnn::writeBUF(buf, o_w);
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
tk::dnn::writeBUF(buf, i_c);
|
||||
tk::dnn::writeBUF(buf, i_h);
|
||||
tk::dnn::writeBUF(buf, i_w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
//todo assert
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "ResizeLayerRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override{
|
||||
ResizeLayerRT *p = new ResizeLayerRT(o_c,o_h,o_w);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
int i_c, i_h, i_w, o_c, o_h, o_w;
|
||||
const char *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
IPluginV2 *clone() const 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 ;
|
||||
|
||||
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(ResizeLayerRTPluginCreator);
|
||||
};
|
||||
|
||||
class ResizeLayerRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
ResizeLayerRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("o_c",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_h",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("o_w",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
ResizeLayerRT *pluginObj = new ResizeLayerRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char *name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
assert(fc->nbFields == 3);
|
||||
assert(fields[0].type == PluginFieldType::kINT32);
|
||||
assert(fields[1].type == PluginFieldType::kINT32);
|
||||
assert(fields[2].type == PluginFieldType::kINT32);
|
||||
int oc = *(static_cast<const int *>(fields[0].data));
|
||||
int oh = *(static_cast<const int *>(fields[1].data));
|
||||
int ow = *(static_cast<const int *>(fields[2].data));
|
||||
ResizeLayerRT *pluginObj = new ResizeLayerRT(oc,oh,ow);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ResizeLayerRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ResizeLayerRTPluginCreator);
|
||||
|
||||
|
||||
@@ -1,181 +1,86 @@
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <vector>
|
||||
#include <NvInfer.h>
|
||||
|
||||
class RouteRT : public IPluginV2 {
|
||||
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) ;
|
||||
|
||||
RouteRT(const void* data,size_t length){
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
groups = readBUF<int>(buf);
|
||||
group_id = readBUF<int>(buf);
|
||||
in = readBUF<int>(buf);
|
||||
for(int i=0;i <MAX_INPUTS;i++){
|
||||
c_in[i] = readBUF<int>(buf);
|
||||
}
|
||||
c= readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override ;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
int out_c = 0;
|
||||
for(int i=0; i<nbInputDims; i++) out_c += inputs[i].d[0];
|
||||
return Dims3{out_c/groups, inputs[0].d[1], inputs[0].d[2]};
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,PluginFormat format, int maxBatchSize) NOEXCEPT override ;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT 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;
|
||||
}
|
||||
int initialize() NOEXCEPT override ;
|
||||
|
||||
int initialize() NOEXCEPT override { return 0;}
|
||||
void terminate() NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override {return 0;}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,cudaStream_t stream) NOEXCEPT override ;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
|
||||
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;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override ;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return (6+MAX_INPUTS)*sizeof(int);
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override ;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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 *getPluginType() const NOEXCEPT override ;
|
||||
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
assert(buf == a + getSerializationSize());
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override ;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "RouteRT_tkDNN";
|
||||
}
|
||||
void destroy() NOEXCEPT override ;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override ;
|
||||
|
||||
void destroy() NOEXCEPT override {delete this; }
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override { return true;}
|
||||
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;
|
||||
};
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override{
|
||||
RouteRT *p = new RouteRT(groups,group_id);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
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);
|
||||
};
|
||||
|
||||
class RouteRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
RouteRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("groups",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("group_id",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char* name,const void* serialData,size_t serialLength) NOEXCEPT override{
|
||||
RouteRT *pluginObj = new RouteRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
assert(fc->nbFields == 2);
|
||||
assert(fields[0].type == PluginFieldType::kINT32);
|
||||
assert(fields[1].type == PluginFieldType::kINT32);
|
||||
int groups = *(static_cast<const int *>(fields[0].data));
|
||||
int group_id = *(static_cast<const int *>(fields[1].data));
|
||||
RouteRT *pluginObj = new RouteRT(groups,group_id);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "RouteRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(RouteRTPluginCreator);
|
||||
|
||||
@@ -1,166 +1,94 @@
|
||||
#ifndef _SHORTCUTRT_PLUGIN_H
|
||||
#define _SHORTCUTRT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
#include <tkdnn.h>
|
||||
|
||||
|
||||
class ShortcutRT : public IPluginV2 {
|
||||
namespace nvinfer1 {
|
||||
|
||||
public:
|
||||
ShortcutRT(tk::dnn::dataDim_t bdim, bool mul) {
|
||||
bDim = bdim;
|
||||
this->bc = bDim.c;
|
||||
this->bh = bDim.h;
|
||||
this->bw = bDim.w;
|
||||
this->mul = mul;
|
||||
}
|
||||
class ShortcutRT : public IPluginV2 {
|
||||
|
||||
~ShortcutRT(){}
|
||||
public:
|
||||
ShortcutRT(tk::dnn::dataDim_t bdim, bool mul);
|
||||
|
||||
ShortcutRT(const void* data,size_t length){
|
||||
const char* buf =reinterpret_cast<const char*>(data),*bufCheck = buf;
|
||||
bDim.c = readBUF<int>(buf);
|
||||
bDim.h = readBUF<int>(buf);
|
||||
bDim.w = readBUF<int>(buf);
|
||||
bDim.l = 1;
|
||||
mul = readBUF<bool>(buf);
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
~ShortcutRT();
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {return 1;}
|
||||
ShortcutRT(const void *data, size_t length);
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3{inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]};
|
||||
}
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT 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() NOEXCEPT override {return 0;}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT 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));
|
||||
shortcutForward(srcDataBack, dstData, batchSize, c, h, w, 1, batchSize, bc, bh, bw, 1, mul, stream);
|
||||
|
||||
return 0;
|
||||
}
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override {
|
||||
return 6*sizeof(int) + sizeof(bool);
|
||||
}
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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, mul);
|
||||
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;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
return true;
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char* getPluginType() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
const char* getPluginVersion() const NOEXCEPT override{
|
||||
return "ShortcutRT_tkDNN";
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
const char* getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override{
|
||||
ShortcutRT *p = new ShortcutRT(bDim,mul);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
IPluginV2 *clone() const NOEXCEPT override;
|
||||
|
||||
int c, h, w;
|
||||
int bc, bh, bw;
|
||||
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;
|
||||
|
||||
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(ShortcutRTPluginCreator);
|
||||
|
||||
int c, h, w;
|
||||
int bc, bh, bw;
|
||||
bool mul;
|
||||
tk::dnn::dataDim_t bDim;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
|
||||
class ShortcutRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
ShortcutRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("bDim",nullptr,PluginFieldType::kUNKNOWN,1));
|
||||
mPluginAttributes.emplace_back(PluginField("mul",nullptr,PluginFieldType::kUNKNOWN,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
ShortcutRT *pluginObj = new ShortcutRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char *name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
//todo assert
|
||||
tk::dnn::dataDim_t bdim = *(static_cast<const tk::dnn::dataDim_t *>(fields[0].data));
|
||||
bool mul = *(static_cast<const bool *>(fields[1].data));
|
||||
ShortcutRT *pluginObj = new ShortcutRT(bdim,mul);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "ShortcutRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
public:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(ShortcutRTPluginCreator);
|
||||
#endif
|
||||
@@ -1,150 +1,87 @@
|
||||
#ifndef _UPSAMPLERT_PLUGIN_H
|
||||
#define _UPSAMPLERT_PLUGIN_H
|
||||
|
||||
#include<cassert>
|
||||
#include "../kernels.h"
|
||||
#include <NvInfer.h>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1 {
|
||||
|
||||
class UpsampleRT : public IPluginV2 {
|
||||
|
||||
public:
|
||||
explicit UpsampleRT(int stride);
|
||||
|
||||
UpsampleRT(const void *data, size_t length);
|
||||
|
||||
~UpsampleRT();
|
||||
|
||||
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;
|
||||
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT override;
|
||||
|
||||
|
||||
class UpsampleRT : public IPluginV2 {
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
public:
|
||||
UpsampleRT(int stride) {
|
||||
this->stride = stride;
|
||||
}
|
||||
void serialize(void *buffer) const NOEXCEPT override;
|
||||
|
||||
UpsampleRT(const void *data,size_t length){
|
||||
const char* buf = reinterpret_cast<const char*>(data),*bufCheck=buf;
|
||||
stride = readBUF<int>(buf);
|
||||
c = readBUF<int>(buf);
|
||||
h = readBUF<int>(buf);
|
||||
w = readBUF<int>(buf);
|
||||
assert(buf == bufCheck + length);
|
||||
}
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override;
|
||||
|
||||
~UpsampleRT(){}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
void destroy() NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) NOEXCEPT override {
|
||||
return Dims3(inputs[0].d[0], inputs[0].d[1]*stride, inputs[0].d[2]*stride);
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
void configureWithFormat (const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs,DataType type,PluginFormat format,int maxBatchSize) NOEXCEPT override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
int initialize() NOEXCEPT override {return 0;}
|
||||
IPluginV2 *clone() const NOEXCEPT override ;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {}
|
||||
int c, h, w, stride;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override { return 0;}
|
||||
class UpsampleRTPluginCreator : public IPluginCreator {
|
||||
public:
|
||||
UpsampleRTPluginCreator();
|
||||
|
||||
virtual int enqueue(int batchSize, const void*const * inputs, void* const* outputs, void* workspace, cudaStream_t stream) NOEXCEPT override {
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT 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;
|
||||
}
|
||||
const char *getPluginNamespace() const NOEXCEPT override;
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
virtual size_t getSerializationSize() const NOEXCEPT override { return 4*sizeof(int);}
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void* buffer) const NOEXCEPT 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());
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
bool supportsFormat(DataType type,PluginFormat format) const NOEXCEPT override{
|
||||
//todo assert
|
||||
return true;
|
||||
}
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
static const char* UPSAMPLE_RT_PLUGIN = "UpsampleRT_TRT";
|
||||
return UPSAMPLE_RT_PLUGIN;
|
||||
}
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
void destroy() NOEXCEPT override{delete this;}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
IPluginV2* clone() const NOEXCEPT override{
|
||||
UpsampleRT *p = new UpsampleRT(stride);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
|
||||
int c, h, w, stride;
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
class UpsampleRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
UpsampleRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("stride",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char* pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char* name,const void* serialData,size_t serialLength) NOEXCEPT override{
|
||||
UpsampleRT *pluginObj = new UpsampleRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
int stride = *(static_cast<const int *>(fields[0].data));
|
||||
UpsampleRT *pluginObj = new UpsampleRT(stride);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
static const char* UPSAMPLE_RT_PLUGIN = "UpsampleRT_TRT";
|
||||
return UPSAMPLE_RT_PLUGIN;
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(UpsampleRTPluginCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(UpsampleRTPluginCreator);
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -1,282 +1,109 @@
|
||||
#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
|
||||
|
||||
namespace nvinfer1 {
|
||||
class YoloRT : public IPluginV2 {
|
||||
|
||||
class YoloRT : public IPluginV2 {
|
||||
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);
|
||||
|
||||
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) {
|
||||
this->yolo = Yolo;
|
||||
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;
|
||||
YoloRT(const void *data, size_t length);
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
YoloRT(const void *data,size_t length){
|
||||
std::vector<float> maskTemp,biasTemp;
|
||||
std::cout<<"LENGTH : "<<length<<std::endl;
|
||||
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);
|
||||
for(int i=0;i<n_masks;i++){
|
||||
maskTemp.push_back(readBUF<dnnType>(buf));
|
||||
std::cout<<maskTemp[i]<<std::endl;
|
||||
}
|
||||
for(int i=0;i<n_masks*2*num;i++){
|
||||
biasTemp.push_back(readBUF<dnnType>(buf));
|
||||
std::cout<<biasTemp[i]<<std::endl;
|
||||
}
|
||||
mask = maskTemp.data();
|
||||
bias = biasTemp.data();
|
||||
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[1] = std::string(tmp);
|
||||
}
|
||||
assert(buf == bufCheck + length);
|
||||
|
||||
}
|
||||
|
||||
~YoloRT() {
|
||||
|
||||
}
|
||||
~YoloRT();
|
||||
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override;
|
||||
|
||||
int getNbOutputs() const NOEXCEPT override {
|
||||
return 1;
|
||||
}
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override;
|
||||
|
||||
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override {
|
||||
return inputs[0];
|
||||
}
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override;
|
||||
|
||||
void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
|
||||
PluginFormat format, int maxBatchSize) NOEXCEPT override {
|
||||
c = inputDims[0].d[0];
|
||||
h = inputDims[0].d[1];
|
||||
w = inputDims[0].d[2];
|
||||
}
|
||||
int initialize() NOEXCEPT override;
|
||||
|
||||
int initialize() NOEXCEPT override {
|
||||
void terminate() NOEXCEPT override;
|
||||
|
||||
return 0;
|
||||
}
|
||||
size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override;
|
||||
|
||||
virtual void terminate() NOEXCEPT override {
|
||||
}
|
||||
|
||||
virtual size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override {
|
||||
return 0;
|
||||
}
|
||||
|
||||
virtual int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) NOEXCEPT 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));
|
||||
int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
||||
cudaStream_t stream) 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
|
||||
size_t getSerializationSize() const NOEXCEPT override;
|
||||
|
||||
if (this->scaleXY != 1)
|
||||
scalAdd(dstData + index, 2 * w * h, this->scaleXY, -0.5 * (this->scaleXY - 1), 1);
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override;
|
||||
|
||||
index = entry_index(b, n * w * h, 4);
|
||||
activationLOGISTICForward(srcData + index, dstData + index, (1 + classes) * w * h, stream);
|
||||
}
|
||||
}
|
||||
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;
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override;
|
||||
|
||||
tk::dnn::Yolo *yolo;
|
||||
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;
|
||||
|
||||
dnnType *mask;
|
||||
dnnType *bias;
|
||||
|
||||
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() const NOEXCEPT 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();
|
||||
|
||||
bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override {
|
||||
return true; //todo implement proper supportsFormat
|
||||
}
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override;
|
||||
|
||||
virtual void serialize(void *buffer) const NOEXCEPT 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;
|
||||
std::cout<<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;
|
||||
}
|
||||
const char *getPluginNamespace() const 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());
|
||||
}
|
||||
IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override;
|
||||
|
||||
const char *getPluginType() const NOEXCEPT override {
|
||||
return "YoloRT_tkDNN";
|
||||
}
|
||||
IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override;
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override {
|
||||
return "1";
|
||||
}
|
||||
const char *getPluginName() const NOEXCEPT override;
|
||||
|
||||
void destroy() NOEXCEPT override { delete this; }
|
||||
const char *getPluginVersion() const NOEXCEPT override;
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override {
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override;
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override {
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
IPluginV2 *clone() const NOEXCEPT override {
|
||||
YoloRT *p = new YoloRT(classes, num,yolo, n_masks, scaleXY, nms_thresh, nms_kind, new_coords);
|
||||
p->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return p;
|
||||
}
|
||||
|
||||
Yolo *yolo;
|
||||
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;
|
||||
|
||||
dnnType *mask;
|
||||
dnnType *bias;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
private:
|
||||
std::string mPluginNamespace;
|
||||
private:
|
||||
static PluginFieldCollection mFC;
|
||||
static std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(YoloRTPluginCreator);
|
||||
};
|
||||
|
||||
class YoloRTPluginCreator : public IPluginCreator{
|
||||
public:
|
||||
YoloRTPluginCreator(){
|
||||
mPluginAttributes.emplace_back(PluginField("classes",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("num",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("yolo",nullptr,PluginFieldType::kUNKNOWN,1));
|
||||
mPluginAttributes.emplace_back(PluginField("numMasks",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("scaleXY",nullptr,PluginFieldType::kFLOAT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("nmsThresh",nullptr,PluginFieldType::kFLOAT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("nmsKind",nullptr,PluginFieldType::kINT32,1));
|
||||
mPluginAttributes.emplace_back(PluginField("newCoords",nullptr,PluginFieldType::kINT32,1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
const char *getPluginNamespace() const NOEXCEPT override{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
IPluginV2 *deserializePlugin(const char *name,const void *serialData,size_t serialLength) NOEXCEPT override{
|
||||
YoloRT *pluginObj = new YoloRT(serialData,serialLength);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
IPluginV2 *createPlugin(const char* name,const PluginFieldCollection *fc) NOEXCEPT override{
|
||||
const PluginField *fields = fc->fields;
|
||||
//todo assert
|
||||
int classes = *(static_cast<const int *>(fields[0].data));
|
||||
int num = *(static_cast<const int *>(fields[1].data));
|
||||
Yolo *yoloTemp = const_cast<Yolo *>(static_cast<const Yolo *>(fields[2].data));
|
||||
int numMasks = *(static_cast<const int*>(fields[3].data));
|
||||
float scaleXY = *(static_cast<const float *>(fields[4].data));
|
||||
float nmsThresh = *(static_cast<const float *>(fields[5].data));
|
||||
int nmsKind = *(static_cast<const int *>(fields[6].data));
|
||||
int newCoords = *(static_cast<const int *>(fields[7].data));
|
||||
YoloRT *pluginObj = new YoloRT(classes,num,yoloTemp,numMasks,scaleXY,nmsThresh,nmsKind,newCoords);
|
||||
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return pluginObj;
|
||||
}
|
||||
|
||||
const char *getPluginName() const NOEXCEPT override{
|
||||
return "YoloRT_tkDNN";
|
||||
}
|
||||
|
||||
const char *getPluginVersion() const NOEXCEPT override{
|
||||
return "1";
|
||||
}
|
||||
|
||||
const PluginFieldCollection *getFieldNames() NOEXCEPT override{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mPluginNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(YoloRTPluginCreator);
|
||||
#endif
|
||||
Reference in New Issue
Block a user