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

TRT8 works with almost every nerual network now!!!!(including demo3d)

Browse Source
This commit is contained in:
perseusdg
2021-10-28 23:35:37 +05:30
parent 8c36dd0431
commit c5e66c6bf6
37 changed files with 1042 additions and 700 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pluginsRT/ReshapeRT.cpp
+64 -24
View File
@@ -4,20 +4,22 @@ using namespace nvinfer1;
std::vector<PluginField> ReshapeRTPluginCreator::mPluginAttributes;
PluginFieldCollection ReshapeRTPluginCreator::mFC{};
ReshapeRT::ReshapeRT(dataDim_t newDim) {
new_dim = newDim;
n = new_dim.n;
c = new_dim.c;
h = new_dim.h;
w = new_dim.w;
static const char* RESHAPERT_PLUGIN_VERSION{"1"};
static const char* RESHAPERT_PLUGIN_NAME{"ReshapeRT_tkDNN"};
ReshapeRT::ReshapeRT(int n,int c,int h,int w) {
this->n = n;
this->c = c;
this->h = h;
this->w = w;
}
ReshapeRT::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);
n = readBUF<int>(buf);
c = readBUF<int>(buf);
h = readBUF<int>(buf);
w = readBUF<int>(buf);
assert(buf == bufCheck + length);
}
@@ -31,8 +33,6 @@ Dims ReshapeRT::getOutputDimensions(int index, const Dims *inputs, int nbInputDi
return Dims3{ c,h,w} ;
}
void ReshapeRT::configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs,DataType type, PluginFormat format, int maxBatchSize) NOEXCEPT {}
int ReshapeRT::initialize() NOEXCEPT {
return 0;
}
@@ -48,11 +48,10 @@ int ReshapeRT::enqueue(int batchSize, const void *const *inputs, void *const *ou
cudaStream_t stream) NOEXCEPT {
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
return 0;
}
#elif NV_TENSORRT_MAJOR == 7
#elif NV_TENSORRT_MAJOR <= 7
int32_t ReshapeRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) {
std::cout << new_dim.c << ":" << new_dim.h << std::endl;
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
@@ -83,11 +82,11 @@ bool ReshapeRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEP
}
const char *ReshapeRT::getPluginType() const NOEXCEPT {
return "ReshapeRT_tkDNN";
return RESHAPERT_PLUGIN_NAME;
}
const char *ReshapeRT::getPluginVersion() const NOEXCEPT {
return "1";
return RESHAPERT_PLUGIN_VERSION;
}
void ReshapeRT::destroy() NOEXCEPT {
@@ -102,14 +101,47 @@ void ReshapeRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
mPluginNamespace = pluginNamespace;
}
IPluginV2 *ReshapeRT::clone() const NOEXCEPT {
auto *p = new ReshapeRT(new_dim);
IPluginV2Ext *ReshapeRT::clone() const NOEXCEPT {
auto *p = new ReshapeRT(n,c,h,w);
p->setPluginNamespace(mPluginNamespace.c_str());
return p;
}
DataType ReshapeRT::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const NOEXCEPT {
return DataType::kFLOAT;
}
void ReshapeRT::attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
IGpuAllocator *gpuAllocator) NOEXCEPT {
}
bool
ReshapeRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted, int nbInputs) const NOEXCEPT {
return false;
}
bool ReshapeRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT {
return false;
}
void ReshapeRT::configurePlugin(const Dims *inputDims, int32_t nbInputs, const Dims *outputDims, int32_t nbOutputs,
const DataType *inputTypes, const DataType *outputTypes, const bool *inputIsBroadcast,
const bool *outputIsBroadcast, PluginFormat floatFormat,
int32_t maxBatchSize) NOEXCEPT {
}
void ReshapeRT::detachFromContext() NOEXCEPT {
}
ReshapeRTPluginCreator::ReshapeRTPluginCreator() {
mPluginAttributes.clear();
mPluginAttributes.emplace_back(PluginField("n", nullptr,PluginFieldType::kINT32,1));
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));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
@@ -122,26 +154,34 @@ const char *ReshapeRTPluginCreator::getPluginNamespace() const NOEXCEPT {
return mPluginNamespace.c_str();
}
IPluginV2 *ReshapeRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT {
IPluginV2Ext *ReshapeRTPluginCreator::deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT {
auto *pluginObj = new ReshapeRT(serialData,serialLength);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
IPluginV2 *ReshapeRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
IPluginV2Ext *ReshapeRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
const PluginField *fields = fc->fields;
dataDim_t newDim = *(static_cast<const dataDim_t *>(fields[0].data));
ReshapeRT *pluginObj = new ReshapeRT(newDim);
assert(fc->nbFields == 4);
for(int i=0;i<4;i++){
assert(fields[1].type == PluginFieldType::kINT32);
}
int n = *(static_cast<const int *>(fields[0].data));
int c = *(static_cast<const int *>(fields[1].data));
int h = *(static_cast<const int *>(fields[2].data));
int w = *(static_cast<const int *>(fields[3].data));
auto *pluginObj = new ReshapeRT(n,c,h,w);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
const char *ReshapeRTPluginCreator::getPluginName() const NOEXCEPT {
return "ReshapeRT_tkDNN";
return RESHAPERT_PLUGIN_NAME;
}
const char *ReshapeRTPluginCreator::getPluginVersion() const NOEXCEPT {
return "1";
return RESHAPERT_PLUGIN_VERSION;
}
const PluginFieldCollection *ReshapeRTPluginCreator::getFieldNames() NOEXCEPT {