362 lines
15 KiB
C++
362 lines
15 KiB
C++
#include <tkDNN/pluginsRT/DeformableConvRT.h>
|
|
|
|
#include <utility>
|
|
using namespace nvinfer1;
|
|
using namespace tk::dnn;
|
|
|
|
std::vector<PluginField> DeformableConvRTPluginCreator::mPluginAttributes;
|
|
PluginFieldCollection DeformableConvRTPluginCreator::mFC{};
|
|
|
|
static const char* DEFORMABLECONVRT_PLUGIN_VERSION{"1"};
|
|
static const char* DEFORMABLECONVRT_PLUGIN_NAME{"DeformableConvRT_tkDNN"};
|
|
|
|
|
|
DeformableConvRT::DeformableConvRT(int chunk_dim, int kh, int kw, int sh, int sw, int ph, int pw, int deformableGroup,
|
|
int i_n, int i_c, int i_h, int i_w, int o_n, int o_c, int o_h, int o_w,std::vector<dnnType> data_H,std::vector<dnnType> bias2_H,
|
|
std::vector<dnnType> ones_d1_h,std::vector<dnnType> ones_d2_h,std::vector<dnnType> offsetH,std::vector<dnnType> maskH,int height_ones,int width_ones,int dim_ones) {
|
|
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->mask_v = std::move(maskH);
|
|
this->offset_v = std::move(offsetH);
|
|
this->ones_d2_v = std::move(ones_d2_h);
|
|
this->ones_d1_v = std::move(ones_d1_h);
|
|
this->data_d_v = std::move(data_H);
|
|
this->bias2_d_v = std::move(bias2_H);
|
|
this->height_ones = height_ones;
|
|
this->width_ones = width_ones;
|
|
this->dim_ones = dim_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(!data_d_v.empty() && !bias2_d_v.empty() && !ones_d1_v.empty() && !ones_d2_v.empty() && !mask_v.empty() && !offset_v.empty()) {
|
|
checkCuda(cudaMemcpy(data_d, data_d_v.data(), sizeof(dnnType) * data_d_v.size(), cudaMemcpyHostToDevice));
|
|
checkCuda(cudaMemcpy(bias2_d, bias2_d_v.data(), sizeof(dnnType) * bias2_d_v.size(), cudaMemcpyHostToDevice));
|
|
checkCuda(cudaMemcpy(ones_d1, ones_d1_v.data(), sizeof(dnnType) * ones_d1_v.size(), cudaMemcpyHostToDevice));
|
|
checkCuda(cudaMemcpy(offset, offset_v.data(), sizeof(dnnType) * offset_v.size(), cudaMemcpyHostToDevice));
|
|
checkCuda(cudaMemcpy(mask, mask_v.data(), sizeof(dnnType) * mask_v.size(), cudaMemcpyHostToDevice));
|
|
checkCuda(cudaMemcpy(ones_d2, ones_d2_v.data(), sizeof(dnnType) * ones_d2_v.size(), cudaMemcpyHostToDevice));
|
|
}
|
|
|
|
|
|
}
|
|
|
|
DeformableConvRT::~DeformableConvRT() {
|
|
checkCuda( cudaFree(data_d) );
|
|
checkCuda( cudaFree(bias2_d) );
|
|
checkCuda( cudaFree(ones_d1) );
|
|
checkCuda( cudaFree(offset) );
|
|
checkCuda( cudaFree(mask) );
|
|
checkCuda( cudaFree(ones_d2) );
|
|
}
|
|
|
|
DeformableConvRT::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);
|
|
height_ones = readBUF<int>(buf);
|
|
width_ones = readBUF<int>(buf);
|
|
dim_ones = readBUF<int>(buf);
|
|
offset_v.resize(chunk_dim*2);
|
|
for(int i=0;i<chunk_dim*2;i++)
|
|
offset_v[i] = readBUF<dnnType>(buf);
|
|
mask_v.resize(chunk_dim);
|
|
for(int i=0;i<chunk_dim;i++)
|
|
mask_v[i] = readBUF<dnnType>(buf);
|
|
data_d_v.resize(i_c*o_c*kh*kw*1);
|
|
for(int i=0;i<(i_c*o_c*kh*kw*1);i++)
|
|
data_d_v[i] = readBUF<dnnType>(buf);
|
|
bias2_d_v.resize(o_c);
|
|
for(int i=0; i < o_c; i++)
|
|
bias2_d_v[i] = readBUF<dnnType>(buf);
|
|
ones_d1_v.resize(height_ones*width_ones);
|
|
for(int i=0; i<height_ones * width_ones; i++)
|
|
ones_d1_v[i] = readBUF<dnnType>(buf);
|
|
ones_d2_v.resize(dim_ones);
|
|
for(int i=0; i<dim_ones; i++)
|
|
ones_d2_v[i] = readBUF<dnnType>(buf);
|
|
assert(buf == bufCheck + length);
|
|
|
|
}
|
|
|
|
int DeformableConvRT::getNbOutputs() const NOEXCEPT {
|
|
return 1;
|
|
}
|
|
|
|
Dims DeformableConvRT::getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT {
|
|
return Dims3{o_c, o_h, o_w};
|
|
}
|
|
|
|
int DeformableConvRT::initialize() NOEXCEPT {
|
|
return 0;
|
|
}
|
|
|
|
void DeformableConvRT::terminate() NOEXCEPT {}
|
|
|
|
size_t DeformableConvRT::getWorkspaceSize(int maxBatchSize) const NOEXCEPT {return 0;}
|
|
|
|
#if NV_TENSORRT_MAJOR > 7
|
|
int DeformableConvRT::enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
|
|
cudaStream_t stream) NOEXCEPT {
|
|
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
|
dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
|
|
|
|
// split conv2d outputs into offset to mask
|
|
for(int b=0; b<batchSize; b++) {
|
|
checkCuda(cudaMemcpy(offset, output_conv + b * 3 * chunk_dim, 2*chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
|
checkCuda(cudaMemcpy(mask, output_conv + b * 3 * chunk_dim + 2*chunk_dim, chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
|
// kernel sigmoid
|
|
activationSIGMOIDForward(mask, mask, chunk_dim);
|
|
// deformable convolution
|
|
dcnV2CudaForward(stat, handle,
|
|
srcData, data_d,
|
|
bias2_d, ones_d1,
|
|
offset, mask,
|
|
reinterpret_cast<dnnType*>(outputs[0]), ones_d2,
|
|
kh, kw,
|
|
sh, sw,
|
|
ph, pw,
|
|
1, 1,
|
|
deformableGroup, b,
|
|
i_n, i_c, i_h, i_w,
|
|
o_n, o_c, o_h, o_w,
|
|
chunk_dim);
|
|
}
|
|
return 0;
|
|
}
|
|
#elif NV_TENSORRT_MAJOR <= 7
|
|
int32_t DeformableConvRT::enqueue(int32_t batchSize, const void *const *inputs, void **outputs, void *workspace,
|
|
cudaStream_t stream) {
|
|
dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
|
|
dnnType *output_conv = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
|
|
|
|
// split conv2d outputs into offset to mask
|
|
for(int b=0; b<batchSize; b++) {
|
|
checkCuda(cudaMemcpy(offset, output_conv + b * 3 * chunk_dim, 2*chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
|
checkCuda(cudaMemcpy(mask, output_conv + b * 3 * chunk_dim + 2*chunk_dim, chunk_dim*sizeof(dnnType), cudaMemcpyDeviceToDevice));
|
|
// kernel sigmoid
|
|
activationSIGMOIDForward(mask, mask, chunk_dim);
|
|
// deformable convolution
|
|
dcnV2CudaForward(stat, handle,
|
|
srcData, data_d,
|
|
bias2_d, ones_d1,
|
|
offset, mask,
|
|
reinterpret_cast<dnnType*>(outputs[0]), ones_d2,
|
|
kh, kw,
|
|
sh, sw,
|
|
ph, pw,
|
|
1, 1,
|
|
deformableGroup, b,
|
|
i_n, i_c, i_h, i_w,
|
|
o_n, o_c, o_h, o_w,
|
|
chunk_dim);
|
|
}
|
|
return 0;
|
|
}
|
|
#endif
|
|
|
|
size_t DeformableConvRT::getSerializationSize() const NOEXCEPT {
|
|
return 19 * 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);
|
|
}
|
|
|
|
void DeformableConvRT::serialize(void *buffer) const NOEXCEPT {
|
|
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
|
|
writeBUF(buf, chunk_dim);
|
|
writeBUF(buf, kh);
|
|
writeBUF(buf, kw);
|
|
writeBUF(buf, sh);
|
|
writeBUF(buf, sw);
|
|
writeBUF(buf, ph);
|
|
writeBUF(buf, pw);
|
|
writeBUF(buf, deformableGroup);
|
|
writeBUF(buf, i_n);
|
|
writeBUF(buf, i_c);
|
|
writeBUF(buf, i_h);
|
|
writeBUF(buf, i_w);
|
|
writeBUF(buf, o_n);
|
|
writeBUF(buf, o_c);
|
|
writeBUF(buf, o_h);
|
|
writeBUF(buf, o_w);
|
|
writeBUF(buf,height_ones);
|
|
writeBUF(buf,width_ones);
|
|
writeBUF(buf,dim_ones);
|
|
for(int i=0; i<offset_v.size(); i++)
|
|
writeBUF(buf, offset_v[i]);
|
|
for(int i=0; i<mask_v.size(); i++)
|
|
writeBUF(buf, mask_v[i]);
|
|
for(int i=0; i<data_d_v.size(); i++)
|
|
writeBUF(buf, data_d_v[i]);
|
|
for(int i=0; i < bias2_d_v.size(); i++)
|
|
writeBUF(buf, bias2_d_v[i]);
|
|
for(int i=0; i<ones_d1_v.size(); i++)
|
|
writeBUF(buf, ones_d1_v[i]);
|
|
for(int i=0; i<ones_d2_v.size(); i++)
|
|
writeBUF(buf, ones_d2_v[i]);
|
|
|
|
assert(buf == a + getSerializationSize());
|
|
}
|
|
|
|
void DeformableConvRT::destroy() NOEXCEPT {
|
|
delete this;
|
|
}
|
|
|
|
|
|
const char *DeformableConvRT::getPluginNamespace() const NOEXCEPT {
|
|
return mPluginNamespace.c_str();
|
|
}
|
|
|
|
void DeformableConvRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
|
|
mPluginNamespace = pluginNamespace;
|
|
}
|
|
|
|
const char *DeformableConvRT::getPluginType() const NOEXCEPT {
|
|
return DEFORMABLECONVRT_PLUGIN_NAME;
|
|
}
|
|
|
|
const char *DeformableConvRT::getPluginVersion() const NOEXCEPT {
|
|
return DEFORMABLECONVRT_PLUGIN_VERSION;
|
|
}
|
|
|
|
IPluginV2Ext *DeformableConvRT::clone() const NOEXCEPT {
|
|
auto *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,data_d_v,bias2_d_v,ones_d1_v,ones_d2_v,offset_v,mask_v,height_ones,width_ones,dim_ones);
|
|
p->setPluginNamespace(mPluginNamespace.c_str());
|
|
return p;
|
|
}
|
|
|
|
DataType
|
|
DeformableConvRT::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const NOEXCEPT {
|
|
return DataType::kFLOAT;
|
|
}
|
|
|
|
void DeformableConvRT::attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
|
|
IGpuAllocator *gpuAllocator) NOEXCEPT {
|
|
handle = cublasContext;
|
|
|
|
}
|
|
|
|
bool DeformableConvRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted,
|
|
int nbInputs) const NOEXCEPT {
|
|
return false;
|
|
}
|
|
|
|
bool DeformableConvRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT {
|
|
return false;
|
|
}
|
|
|
|
void DeformableConvRT::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 DeformableConvRT::detachFromContext() NOEXCEPT {
|
|
|
|
}
|
|
|
|
bool DeformableConvRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
|
|
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
|
}
|
|
|
|
DeformableConvRTPluginCreator::DeformableConvRTPluginCreator() {
|
|
mPluginAttributes.clear();
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
void DeformableConvRTPluginCreator::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
|
|
mPluginNamespace = pluginNamespace;
|
|
}
|
|
|
|
const char *DeformableConvRTPluginCreator::getPluginNamespace() const NOEXCEPT {
|
|
return mPluginNamespace.c_str();
|
|
}
|
|
|
|
IPluginV2Ext *DeformableConvRTPluginCreator::deserializePlugin(const char *name, const void *serialData,
|
|
size_t serialLength) NOEXCEPT {
|
|
auto *pluginObj = new DeformableConvRT(serialData,serialLength);
|
|
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
|
return pluginObj;
|
|
}
|
|
|
|
IPluginV2Ext *DeformableConvRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
|
|
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[15].data));
|
|
std::vector<dnnType> mask_v(static_cast<const dnnType*>(fields[16].data),static_cast<const dnnType*>(fields[16].data)+fields[16].length);
|
|
std::vector<dnnType> offset_v(static_cast<const dnnType*>(fields[17].data),static_cast<const dnnType*>(fields[17].data)+fields[17].length);
|
|
std::vector<dnnType> ones_d2_v(static_cast<const dnnType*>(fields[18].data),static_cast<const dnnType*>(fields[18].data)+fields[18].length);
|
|
std::vector<dnnType> ones_d1_v(static_cast<const dnnType*>(fields[19].data),static_cast<const dnnType*>(fields[19].data)+fields[19].length);
|
|
std::vector<dnnType> data_d_v(static_cast<const dnnType*>(fields[20].data),static_cast<const dnnType*>(fields[20].data)+fields[20].length);
|
|
std::vector<dnnType> bias2_d_v(static_cast<const dnnType*>(fields[21].data),static_cast<const dnnType*>(fields[21].data)+fields[21].length);
|
|
int height_ones = *(static_cast<const int *>(fields[22].data));
|
|
int width_ones = *(static_cast<const int *>(fields[23].data));
|
|
int dim_ones = *(static_cast<const int *>(fields[24].data));
|
|
auto *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,data_d_v,bias2_d_v,ones_d1_v,ones_d2_v,offset_v,mask_v,height_ones,width_ones,dim_ones);
|
|
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
|
|
return pluginObj;
|
|
}
|
|
|
|
const char *DeformableConvRTPluginCreator::getPluginName() const NOEXCEPT {
|
|
return DEFORMABLECONVRT_PLUGIN_NAME;
|
|
}
|
|
|
|
const char *DeformableConvRTPluginCreator::getPluginVersion() const NOEXCEPT {
|
|
return DEFORMABLECONVRT_PLUGIN_VERSION;
|
|
}
|
|
|
|
const PluginFieldCollection *DeformableConvRTPluginCreator::getFieldNames() NOEXCEPT {
|
|
return &mFC;
|
|
}
|
|
|
|
|
|
|
|
|
|
|