Unverified Commit 2e5628b4 authored by q.yao's avatar q.yao Committed by GitHub
Browse files

[Refactor]: Remove deployment for dev-2.x (#2225)

* remove deploy for 2.0

* update onnx ut
parent 961373ad
// Copyright (c) OpenMMLab. All rights reserved
#include "trt_plugin.hpp"
#include "trt_corner_pool.hpp"
#include "trt_cummaxmin.hpp"
#include "trt_deform_conv.hpp"
#include "trt_grid_sampler.hpp"
#include "trt_instance_norm.hpp"
#include "trt_modulated_deform_conv.hpp"
#include "trt_nms.hpp"
#include "trt_roi_align.hpp"
#include "trt_scatternd.hpp"
REGISTER_TENSORRT_PLUGIN(CumMaxPluginDynamicCreator);
REGISTER_TENSORRT_PLUGIN(CumMinPluginDynamicCreator);
REGISTER_TENSORRT_PLUGIN(GridSamplerDynamicCreator);
REGISTER_TENSORRT_PLUGIN(DeformableConvPluginDynamicCreator);
REGISTER_TENSORRT_PLUGIN(ModulatedDeformableConvPluginDynamicCreator);
REGISTER_TENSORRT_PLUGIN(NonMaxSuppressionDynamicCreator);
REGISTER_TENSORRT_PLUGIN(RoIAlignPluginDynamicCreator);
REGISTER_TENSORRT_PLUGIN(ONNXScatterNDDynamicCreator);
REGISTER_TENSORRT_PLUGIN(InstanceNormalizationDynamicCreator);
REGISTER_TENSORRT_PLUGIN(CornerPoolPluginDynamicCreator);
extern "C" {
bool initLibMMCVInferPlugins() { return true; }
} // extern "C"
// Copyright (c) OpenMMLab. All rights reserved
#include "trt_roi_align.hpp"
#include <assert.h>
#include <chrono>
#include "trt_serialize.hpp"
extern void TRTRoIAlignForwardCUDAKernelLauncher_float(
const float *input, const float *rois, float *output, float *argmax_y,
float *argmax_x, int output_size, int channels, int height, int width,
int aligned_height, int aligned_width, float spatial_scale,
int sampling_ratio, int pool_mode, bool aligned, cudaStream_t stream);
namespace {
static const char *PLUGIN_VERSION{"1"};
static const char *PLUGIN_NAME{"MMCVRoiAlign"};
} // namespace
nvinfer1::PluginFieldCollection RoIAlignPluginDynamicCreator::mFC{};
std::vector<nvinfer1::PluginField>
RoIAlignPluginDynamicCreator::mPluginAttributes;
RoIAlignPluginDynamic::RoIAlignPluginDynamic(const std::string &name,
int outWidth, int outHeight,
float spatialScale,
int sampleRatio, int poolMode,
bool aligned)
: mLayerName(name),
mOutWidth(outWidth),
mOutHeight(outHeight),
mSpatialScale(spatialScale),
mSampleRatio(sampleRatio),
mPoolMode(poolMode),
mAligned(aligned) {}
RoIAlignPluginDynamic::RoIAlignPluginDynamic(const std::string name,
const void *data, size_t length)
: mLayerName(name) {
deserialize_value(&data, &length, &mOutWidth);
deserialize_value(&data, &length, &mOutHeight);
deserialize_value(&data, &length, &mSpatialScale);
deserialize_value(&data, &length, &mSampleRatio);
deserialize_value(&data, &length, &mPoolMode);
deserialize_value(&data, &length, &mAligned);
}
nvinfer1::IPluginV2DynamicExt *RoIAlignPluginDynamic::clone() const {
RoIAlignPluginDynamic *plugin = new RoIAlignPluginDynamic(
mLayerName, mOutWidth, mOutHeight, mSpatialScale, mSampleRatio, mPoolMode,
mAligned);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
nvinfer1::DimsExprs RoIAlignPluginDynamic::getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) {
nvinfer1::DimsExprs ret;
ret.nbDims = 4;
ret.d[0] = inputs[1].d[0];
ret.d[1] = inputs[0].d[1];
ret.d[2] = exprBuilder.constant(mOutHeight);
ret.d[3] = exprBuilder.constant(mOutWidth);
return ret;
}
bool RoIAlignPluginDynamic::supportsFormatCombination(
int pos, const nvinfer1::PluginTensorDesc *inOut, int nbInputs,
int nbOutputs) {
return inOut[pos].type == nvinfer1::DataType::kFLOAT &&
inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
}
void RoIAlignPluginDynamic::configurePlugin(
const nvinfer1::DynamicPluginTensorDesc *inputs, int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *outputs, int nbOutputs) {}
size_t RoIAlignPluginDynamic::getWorkspaceSize(
const nvinfer1::PluginTensorDesc *inputs, int nbInputs,
const nvinfer1::PluginTensorDesc *outputs, int nbOutputs) const {
size_t output_size = 0;
size_t word_size = 0;
switch (mPoolMode) {
case 0: // max
output_size = outputs[0].dims.d[0] * outputs[0].dims.d[1] *
outputs[0].dims.d[2] * outputs[0].dims.d[3];
word_size = mmcv::getElementSize(outputs[0].type);
return output_size * word_size * 2;
break;
case 1:
return 0;
break;
default:
return 0;
}
return 0;
}
int RoIAlignPluginDynamic::enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs,
void *const *outputs, void *workSpace,
cudaStream_t stream) {
int channels = inputDesc[0].dims.d[1];
int height = inputDesc[0].dims.d[2];
int width = inputDesc[0].dims.d[3];
int output_size = outputDesc[0].dims.d[0] * outputDesc[0].dims.d[1] *
outputDesc[0].dims.d[2] * outputDesc[0].dims.d[3];
int word_size = mmcv::getElementSize(outputDesc[0].type);
const void *feat = inputs[0];
const void *rois = inputs[1];
void *output = outputs[0];
void *argmax_y = nullptr;
void *argmax_x = nullptr;
switch (mPoolMode) {
case 0: // max
argmax_y = workSpace;
argmax_x = argmax_y + output_size * word_size;
break;
case 1: // avg
break;
}
switch (outputDesc[0].type) {
case nvinfer1::DataType::kFLOAT:
TRTRoIAlignForwardCUDAKernelLauncher_float(
(const float *)feat, (const float *)rois, (float *)output,
(float *)argmax_y, (float *)argmax_x, output_size, channels, height,
width, mOutHeight, mOutWidth, mSpatialScale, mSampleRatio, mPoolMode,
mAligned, stream);
break;
default:
break;
}
return 0;
}
nvinfer1::DataType RoIAlignPluginDynamic::getOutputDataType(
int index, const nvinfer1::DataType *inputTypes, int nbInputs) const {
return inputTypes[0];
}
// IPluginV2 Methods
const char *RoIAlignPluginDynamic::getPluginType() const { return PLUGIN_NAME; }
const char *RoIAlignPluginDynamic::getPluginVersion() const {
return PLUGIN_VERSION;
}
int RoIAlignPluginDynamic::getNbOutputs() const { return 1; }
int RoIAlignPluginDynamic::initialize() { return 0; }
void RoIAlignPluginDynamic::terminate() {}
size_t RoIAlignPluginDynamic::getSerializationSize() const {
return sizeof(mOutWidth) + sizeof(mOutHeight) + sizeof(mSpatialScale) +
sizeof(mSampleRatio) + sizeof(mPoolMode) + sizeof(mAligned);
}
void RoIAlignPluginDynamic::serialize(void *buffer) const {
serialize_value(&buffer, mOutWidth);
serialize_value(&buffer, mOutHeight);
serialize_value(&buffer, mSpatialScale);
serialize_value(&buffer, mSampleRatio);
serialize_value(&buffer, mPoolMode);
serialize_value(&buffer, mAligned);
}
void RoIAlignPluginDynamic::destroy() {
// This gets called when the network containing plugin is destroyed
delete this;
}
void RoIAlignPluginDynamic::setPluginNamespace(const char *libNamespace) {
mNamespace = libNamespace;
}
const char *RoIAlignPluginDynamic::getPluginNamespace() const {
return mNamespace.c_str();
}
////////////////////// creator /////////////////////////////
RoIAlignPluginDynamicCreator::RoIAlignPluginDynamicCreator() {
mPluginAttributes.emplace_back(nvinfer1::PluginField("output_height"));
mPluginAttributes.emplace_back(nvinfer1::PluginField("output_width"));
mPluginAttributes.emplace_back(nvinfer1::PluginField("spatial_scale"));
mPluginAttributes.emplace_back(nvinfer1::PluginField("sampling_ratio"));
mPluginAttributes.emplace_back(nvinfer1::PluginField("mode"));
mPluginAttributes.emplace_back(nvinfer1::PluginField("aligned"));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
const char *RoIAlignPluginDynamicCreator::getPluginName() const {
return PLUGIN_NAME;
}
const char *RoIAlignPluginDynamicCreator::getPluginVersion() const {
return PLUGIN_VERSION;
}
const nvinfer1::PluginFieldCollection *
RoIAlignPluginDynamicCreator::getFieldNames() {
return &mFC;
}
nvinfer1::IPluginV2 *RoIAlignPluginDynamicCreator::createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) {
int outWidth = 7;
int outHeight = 7;
float spatialScale = 1.0;
int sampleRatio = 0;
int poolMode = -1;
bool aligned = true;
for (int i = 0; i < fc->nbFields; i++) {
if (fc->fields[i].data == nullptr) {
continue;
}
std::string field_name(fc->fields[i].name);
if (field_name.compare("output_height") == 0) {
outHeight = static_cast<const int *>(fc->fields[i].data)[0];
}
if (field_name.compare("output_width") == 0) {
outWidth = static_cast<const int *>(fc->fields[i].data)[0];
}
if (field_name.compare("spatial_scale") == 0) {
spatialScale = static_cast<const float *>(fc->fields[i].data)[0];
}
if (field_name.compare("sampling_ratio") == 0) {
sampleRatio = static_cast<const int *>(fc->fields[i].data)[0];
}
if (field_name.compare("mode") == 0) {
int data_size = fc->fields[i].length;
const char *data_start = static_cast<const char *>(fc->fields[i].data);
std::string poolModeStr(data_start, data_size);
if (poolModeStr == "avg") {
poolMode = 1;
} else if (poolModeStr == "max") {
poolMode = 0;
} else {
std::cout << "Unknown pool mode \"" << poolModeStr << "\"."
<< std::endl;
}
assert(poolMode >= 0);
}
if (field_name.compare("aligned") == 0) {
int aligned_int = static_cast<const int *>(fc->fields[i].data)[0];
aligned = aligned_int != 0;
}
}
assert(outHeight > 0);
assert(outWidth > 0);
assert(spatialScale > 0.);
assert(poolMode >= 0);
RoIAlignPluginDynamic *plugin = new RoIAlignPluginDynamic(
name, outWidth, outHeight, spatialScale, sampleRatio, poolMode, aligned);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
nvinfer1::IPluginV2 *RoIAlignPluginDynamicCreator::deserializePlugin(
const char *name, const void *serialData, size_t serialLength) {
auto plugin = new RoIAlignPluginDynamic(name, serialData, serialLength);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
void RoIAlignPluginDynamicCreator::setPluginNamespace(
const char *libNamespace) {
mNamespace = libNamespace;
}
const char *RoIAlignPluginDynamicCreator::getPluginNamespace() const {
return mNamespace.c_str();
}
// Copyright (c) OpenMMLab. All rights reserved
#include "common_cuda_helper.hpp"
#include "roi_align_cuda_kernel.cuh"
template <typename scalar_t>
void TRTRoIAlignForwardCUDAKernelLauncher(
const scalar_t* input, const scalar_t* rois, scalar_t* output,
scalar_t* argmax_y, scalar_t* argmax_x, int output_size, int channels,
int height, int width, int aligned_height, int aligned_width,
scalar_t spatial_scale, int sampling_ratio, int pool_mode, bool aligned,
cudaStream_t stream) {
roi_align_forward_cuda_kernel<scalar_t>
<<<GET_BLOCKS(output_size), THREADS_PER_BLOCK, 0, stream>>>(
output_size, input, rois, output, argmax_y, argmax_x, aligned_height,
aligned_width, static_cast<scalar_t>(spatial_scale), sampling_ratio,
pool_mode, aligned, channels, height, width);
}
void TRTRoIAlignForwardCUDAKernelLauncher_float(
const float* input, const float* rois, float* output, float* argmax_y,
float* argmax_x, int output_size, int channels, int height, int width,
int aligned_height, int aligned_width, float spatial_scale,
int sampling_ratio, int pool_mode, bool aligned, cudaStream_t stream) {
TRTRoIAlignForwardCUDAKernelLauncher<float>(
input, rois, output, argmax_y, argmax_x, output_size, channels, height,
width, aligned_height, aligned_width, spatial_scale, sampling_ratio,
pool_mode, aligned, stream);
}
// Copyright (c) OpenMMLab. All rights reserved
#include "trt_scatternd.hpp"
#include <assert.h>
#include <stdio.h>
#include <chrono>
#include "trt_serialize.hpp"
extern void TRTONNXScatterNDKernelLauncher_float(
const float *data, const int *indices, const float *update, const int *dims,
int nbDims, const int *indices_dims, int indice_nbDims, float *output,
cudaStream_t stream);
extern void TRTONNXScatterNDKernelLauncher_int32(
const int *data, const int *indices, const int *update, const int *dims,
int nbDims, const int *indices_dims, int indice_nbDims, int *output,
cudaStream_t stream);
namespace {
static const char *PLUGIN_VERSION{"1"};
static const char *PLUGIN_NAME{"ScatterND"};
} // namespace
nvinfer1::PluginFieldCollection ONNXScatterNDDynamicCreator::mFC{};
std::vector<nvinfer1::PluginField>
ONNXScatterNDDynamicCreator::mPluginAttributes;
ONNXScatterNDDynamic::ONNXScatterNDDynamic(const std::string &name)
: mLayerName(name) {}
ONNXScatterNDDynamic::ONNXScatterNDDynamic(const std::string name,
const void *data, size_t length)
: mLayerName(name) {}
nvinfer1::IPluginV2DynamicExt *ONNXScatterNDDynamic::clone() const {
ONNXScatterNDDynamic *plugin = new ONNXScatterNDDynamic(mLayerName);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
nvinfer1::DimsExprs ONNXScatterNDDynamic::getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) {
return inputs[0];
}
bool ONNXScatterNDDynamic::supportsFormatCombination(
int pos, const nvinfer1::PluginTensorDesc *inOut, int nbInputs,
int nbOutputs) {
if (pos < nbInputs) {
switch (pos) {
case 0:
// data
return (inOut[pos].type == nvinfer1::DataType::kFLOAT &&
inOut[pos].format == nvinfer1::TensorFormat::kLINEAR) ||
(inOut[pos].type == nvinfer1::DataType::kINT32 &&
inOut[pos].format == nvinfer1::TensorFormat::kLINEAR);
case 1:
// indices
return inOut[pos].type == nvinfer1::DataType::kINT32 &&
inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
case 2:
// updates
return inOut[pos].type == inOut[0].type &&
inOut[pos].format == inOut[0].format;
default:
return true;
}
} else {
switch (pos - nbInputs) {
case 0:
// output
return inOut[pos].type == inOut[0].type &&
inOut[pos].format == inOut[0].format;
default:
return true;
}
}
return true;
}
void ONNXScatterNDDynamic::configurePlugin(
const nvinfer1::DynamicPluginTensorDesc *inputs, int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *outputs, int nbOutputs) {}
size_t ONNXScatterNDDynamic::getWorkspaceSize(
const nvinfer1::PluginTensorDesc *inputs, int nbInputs,
const nvinfer1::PluginTensorDesc *outputs, int nbOutputs) const {
return 0;
}
int ONNXScatterNDDynamic::enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs,
void *const *outputs, void *workSpace,
cudaStream_t stream) {
const int *dims = &(inputDesc[0].dims.d[0]);
const int *indices_dims = &(inputDesc[1].dims.d[0]);
int nbDims = inputDesc[0].dims.nbDims;
int indice_nbDims = inputDesc[1].dims.nbDims;
const void *data = inputs[0];
const void *indices = inputs[1];
const void *update = inputs[2];
void *output = outputs[0];
auto data_type = inputDesc[0].type;
switch (data_type) {
case nvinfer1::DataType::kFLOAT:
TRTONNXScatterNDKernelLauncher_float(
(float *)data, (int *)indices, (float *)update, dims, nbDims,
indices_dims, indice_nbDims, (float *)output, stream);
break;
case nvinfer1::DataType::kINT32:
TRTONNXScatterNDKernelLauncher_int32(
(int *)data, (int *)indices, (int *)update, dims, nbDims,
indices_dims, indice_nbDims, (int *)output, stream);
break;
default:
break;
}
return 0;
}
nvinfer1::DataType ONNXScatterNDDynamic::getOutputDataType(
int index, const nvinfer1::DataType *inputTypes, int nbInputs) const {
return inputTypes[0];
}
// IPluginV2 Methods
const char *ONNXScatterNDDynamic::getPluginType() const { return PLUGIN_NAME; }
const char *ONNXScatterNDDynamic::getPluginVersion() const {
return PLUGIN_VERSION;
}
int ONNXScatterNDDynamic::getNbOutputs() const { return 1; }
int ONNXScatterNDDynamic::initialize() { return 0; }
void ONNXScatterNDDynamic::terminate() {}
size_t ONNXScatterNDDynamic::getSerializationSize() const { return 0; }
void ONNXScatterNDDynamic::serialize(void *buffer) const {}
void ONNXScatterNDDynamic::destroy() {
// This gets called when the network containing plugin is destroyed
delete this;
}
void ONNXScatterNDDynamic::setPluginNamespace(const char *libNamespace) {
mNamespace = libNamespace;
}
const char *ONNXScatterNDDynamic::getPluginNamespace() const {
return mNamespace.c_str();
}
////////////////////// creator /////////////////////////////
ONNXScatterNDDynamicCreator::ONNXScatterNDDynamicCreator() {
mPluginAttributes.clear();
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
const char *ONNXScatterNDDynamicCreator::getPluginName() const {
return PLUGIN_NAME;
}
const char *ONNXScatterNDDynamicCreator::getPluginVersion() const {
return PLUGIN_VERSION;
}
const nvinfer1::PluginFieldCollection *
ONNXScatterNDDynamicCreator::getFieldNames() {
return &mFC;
}
nvinfer1::IPluginV2 *ONNXScatterNDDynamicCreator::createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) {
ONNXScatterNDDynamic *plugin = new ONNXScatterNDDynamic(name);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
nvinfer1::IPluginV2 *ONNXScatterNDDynamicCreator::deserializePlugin(
const char *name, const void *serialData, size_t serialLength) {
auto plugin = new ONNXScatterNDDynamic(name, serialData, serialLength);
plugin->setPluginNamespace(getPluginNamespace());
return plugin;
}
void ONNXScatterNDDynamicCreator::setPluginNamespace(const char *libNamespace) {
mNamespace = libNamespace;
}
const char *ONNXScatterNDDynamicCreator::getPluginNamespace() const {
return mNamespace.c_str();
}
// Copyright (c) OpenMMLab. All rights reserved
#include <stdio.h>
#include <vector>
#include "common_cuda_helper.hpp"
#include "trt_cuda_helper.cuh"
#include "trt_plugin_helper.hpp"
static int const threadsPerBlock = sizeof(unsigned long long int) * 8;
using mmcv::TensorDesc;
template <typename T>
__global__ void onnx_scatternd_kernel(const int n, const int* indices,
const T* update, T* output,
TensorDesc tensor_desc,
TensorDesc indice_desc) {
const int indice_cols = indice_desc.shape[indice_desc.dim - 1];
const int copy_stride = tensor_desc.stride[indice_cols - 1];
const int* stride = &(tensor_desc.stride[0]);
CUDA_1D_KERNEL_LOOP(index, n) {
int output_offset = 0;
const int* indices_current = indices + index * indice_cols;
for (int i = 0; i < indice_cols; ++i) {
output_offset += stride[i] * indices_current[i];
}
memcpy(output + output_offset, update + index * copy_stride,
copy_stride * sizeof(T));
}
}
template <typename T>
void TRTONNXScatterNDKernelLauncher(const T* data, const int* indices,
const T* update, const int* dims,
int nbDims, const int* indices_dims,
int indice_nbDims, T* output,
cudaStream_t stream) {
// fill tensordesc and initial
TensorDesc tensor_desc;
memset((void*)&tensor_desc, 0, sizeof(TensorDesc));
tensor_desc.dim = nbDims;
tensor_desc.shape[nbDims - 1] = dims[nbDims - 1];
tensor_desc.stride[nbDims - 1] = 1;
for (int i = nbDims - 2; i >= 0; --i) {
tensor_desc.shape[i] = dims[i];
tensor_desc.stride[i] = dims[i + 1] * tensor_desc.stride[i + 1];
}
const int data_size = tensor_desc.stride[0] * tensor_desc.shape[0];
TensorDesc indice_desc;
memset((void*)&indice_desc, 0, sizeof(TensorDesc));
indice_desc.dim = indice_nbDims;
indice_desc.shape[indice_nbDims - 1] = indices_dims[indice_nbDims - 1];
indice_desc.stride[indice_nbDims - 1] = 1;
for (int i = indice_nbDims - 2; i >= 0; --i) {
indice_desc.shape[i] = indices_dims[i];
indice_desc.stride[i] = indices_dims[i + 1] * indice_desc.stride[i + 1];
}
// output = np.copy(data)
cudaMemcpyAsync(output, data, data_size * sizeof(T),
cudaMemcpyDeviceToDevice);
int num_update_indice = 1;
for (int i = 0; i < indice_nbDims - 1; ++i) {
num_update_indice *= indice_desc.shape[i];
}
// scatter
const int col_block = GET_BLOCKS(num_update_indice, threadsPerBlock);
onnx_scatternd_kernel<<<col_block, threadsPerBlock, 0, stream>>>(
num_update_indice, indices, update, output, tensor_desc, indice_desc);
}
void TRTONNXScatterNDKernelLauncher_float(const float* data, const int* indices,
const float* update, const int* dims,
int nbDims, const int* indices_dims,
int indice_nbDims, float* output,
cudaStream_t stream) {
TRTONNXScatterNDKernelLauncher<float>(data, indices, update, dims, nbDims,
indices_dims, indice_nbDims, output,
stream);
}
void TRTONNXScatterNDKernelLauncher_int32(const int* data, const int* indices,
const int* update, const int* dims,
int nbDims, const int* indices_dims,
int indice_nbDims, int* output,
cudaStream_t stream) {
TRTONNXScatterNDKernelLauncher<int>(data, indices, update, dims, nbDims,
indices_dims, indice_nbDims, output,
stream);
}
#ifndef TRT_CORNER_POOL_HPP
#define TRT_CORNER_POOL_HPP
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
enum TRT_CORNER_POOL_TYPE {
TRT_TOP_POOL = 0,
TRT_BOTTOM_POOL = 1,
TRT_LEFT_POOL = 2,
TRT_RIGHT_POOL = 3
};
// implement of CornerPool
class CornerPoolPluginDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
CornerPoolPluginDynamic(const std::string &name,
TRT_CORNER_POOL_TYPE poolType);
CornerPoolPluginDynamic(const std::string name, const void *data,
size_t length);
CornerPoolPluginDynamic() = delete;
~CornerPoolPluginDynamic();
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
protected:
const std::string mLayerName;
std::string mNamespace;
TRT_CORNER_POOL_TYPE mPoolType;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
// CornerPool creator
class CornerPoolPluginDynamicCreator : public nvinfer1::IPluginCreator {
public:
CornerPoolPluginDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
protected:
nvinfer1::PluginFieldCollection mFC;
std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif TRT_CORNER_POOL_HPP // TRT_CORNER_POOL_HPP
// Copyright (c) OpenMMLab. All rights reserved
#ifndef TRT_CUDA_HELPER_HPP
#define TRT_CUDA_HELPER_HPP
#include <cublas_v2.h>
#define cudaCheckError() \
{ \
cudaError_t e = cudaGetLastError(); \
if (e != cudaSuccess) { \
printf("Cuda failure %s:%d: '%s'\n", __FILE__, __LINE__, \
cudaGetErrorString(e)); \
exit(0); \
} \
}
/**
* Returns a view of the original tensor with its dimensions permuted.
*
* @param[out] dst pointer to the destination tensor
* @param[in] src pointer to the source tensor
* @param[in] src_size shape of the src tensor
* @param[in] permute The desired ordering of dimensions
* @param[in] src_dim dim of src tensor
* @param[in] stream cuda stream handle
*/
template <class scalar_t>
void memcpyPermute(scalar_t* dst, const scalar_t* src, int* src_size,
int* permute, int src_dim, cudaStream_t stream = 0);
template <typename scalar_t>
cublasStatus_t cublasGemmWrap(cublasHandle_t handle, cublasOperation_t transa,
cublasOperation_t transb, int m, int n, int k,
const scalar_t* alpha, const scalar_t* A, int lda,
const scalar_t* B, int ldb, const scalar_t* beta,
scalar_t* C, int ldc) {
return CUBLAS_STATUS_INTERNAL_ERROR;
}
#endif // TRT_CUDA_HELPER_HPP
#ifndef TRT_CUMMAXMIN_HPP
#define TRT_CUMMAXMIN_HPP
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
enum TRT_CUMCMPTYPE { TRT_CUMMAX = 0, TRT_CUMMIN = 1 };
// implement of cummax and cummin
class CumMaxMinPluginDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
CumMaxMinPluginDynamic(const std::string &name, int dim,
TRT_CUMCMPTYPE cumType);
CumMaxMinPluginDynamic(const std::string name, const void *data,
size_t length);
CumMaxMinPluginDynamic() = delete;
~CumMaxMinPluginDynamic();
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
protected:
const std::string mLayerName;
std::string mNamespace;
int mDim;
TRT_CUMCMPTYPE mCumType;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
// cummax and cummin creator
class CumMaxMinPluginDynamicCreator : public nvinfer1::IPluginCreator {
public:
CumMaxMinPluginDynamicCreator(TRT_CUMCMPTYPE cumType);
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
protected:
TRT_CUMCMPTYPE mCumType;
nvinfer1::PluginFieldCollection mFC;
std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
// cummax creator
class CumMaxPluginDynamicCreator : public CumMaxMinPluginDynamicCreator {
public:
CumMaxPluginDynamicCreator();
const char *getPluginName() const override;
};
// cummin creator
class CumMinPluginDynamicCreator : public CumMaxMinPluginDynamicCreator {
public:
CumMinPluginDynamicCreator();
const char *getPluginName() const override;
};
#endif TRT_CUMMAXMIN_HPP // TRT_CUMMAXMIN_HPP
#ifndef TRT_DEFORM_CONV_HPP
#define TRT_DEFORM_CONV_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
class DeformableConvPluginDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
DeformableConvPluginDynamic(const std::string &name,
const nvinfer1::Dims &stride,
const nvinfer1::Dims &padding,
const nvinfer1::Dims &dilation,
const int deformableGroup, const int group,
int im2colStep);
DeformableConvPluginDynamic(const std::string name, const void *data,
size_t length);
DeformableConvPluginDynamic() = delete;
~DeformableConvPluginDynamic();
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
void attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
nvinfer1::IGpuAllocator *gpuAllocator) override;
void detachFromContext() override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
nvinfer1::Dims mStride;
nvinfer1::Dims mPadding;
nvinfer1::Dims mDilation;
int mDeformableGroup;
int mGroup;
int mIm2colStep;
cublasHandle_t m_cublas_handle;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class DeformableConvPluginDynamicCreator : public nvinfer1::IPluginCreator {
public:
DeformableConvPluginDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_DEFORM_CONV_HPP
#ifndef TRT_GRID_SAMPLER_HPP
#define TRT_GRID_SAMPLER_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
namespace mmcv {
enum class GridSamplerInterpolation { Bilinear, Nearest };
enum class GridSamplerPadding { Zeros, Border, Reflection };
} // namespace mmcv
class GridSamplerDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
GridSamplerDynamic(const std::string &name, int mode, int paddingMode,
bool alignCorners);
GridSamplerDynamic(const std::string name, const void *data, size_t length);
GridSamplerDynamic() = delete;
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
int mMode;
int mPaddingMode;
bool mAlignCorners;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class GridSamplerDynamicCreator : public nvinfer1::IPluginCreator {
public:
GridSamplerDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_GRID_SAMPLER_HPP
// Modified from:
// https://github.com/NVIDIA/TensorRT/blob/master/plugin/instanceNormalizationPlugin/instanceNormalizationPlugin.h
#ifndef TRT_INSTANCE_NORMALIZATION_PLUGIN_H
#define TRT_INSTANCE_NORMALIZATION_PLUGIN_H
#include <cudnn.h>
#include <iostream>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
typedef unsigned short half_type;
class InstanceNormalizationDynamic final
: public nvinfer1::IPluginV2DynamicExt {
public:
InstanceNormalizationDynamic(const std::string& name, float epsilon);
InstanceNormalizationDynamic(const std::string& name, void const* serialData,
size_t serialLength);
InstanceNormalizationDynamic() = delete;
~InstanceNormalizationDynamic() override;
int getNbOutputs() const override;
// DynamicExt plugins returns DimsExprs class instead of Dims
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs,
nvinfer1::IExprBuilder& exprBuilder) override;
int initialize() override;
void terminate() override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc* outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
const nvinfer1::PluginTensorDesc* outputDesc,
const void* const* inputs, void* const* outputs, void* workspace,
cudaStream_t stream) override;
size_t getSerializationSize() const override;
void serialize(void* buffer) const override;
// DynamicExt plugin supportsFormat update.
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc* inOut,
int nbInputs, int nbOutputs) override;
const char* getPluginType() const override;
const char* getPluginVersion() const override;
void destroy() override;
nvinfer1::IPluginV2DynamicExt* clone() const override;
void setPluginNamespace(const char* pluginNamespace) override;
const char* getPluginNamespace() const override;
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType* inputTypes,
int nbInputs) const override;
void attachToContext(cudnnContext* cudnn, cublasContext* cublas,
nvinfer1::IGpuAllocator* allocator) override;
void detachFromContext() override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc* out,
int nbOutputs) override;
private:
const std::string mLayerName;
float mEpsilon{};
cudnnHandle_t _cudnn_handle{};
cudnnTensorDescriptor_t _x_desc{}, _y_desc{}, _b_desc{};
std::string mPluginNamespace{};
};
class InstanceNormalizationDynamicCreator : public nvinfer1::IPluginCreator {
public:
InstanceNormalizationDynamicCreator();
~InstanceNormalizationDynamicCreator() override = default;
const char* getPluginName() const override;
const char* getPluginVersion() const override;
const nvinfer1::PluginFieldCollection* getFieldNames() override;
nvinfer1::IPluginV2DynamicExt* createPlugin(
const char* name, const nvinfer1::PluginFieldCollection* fc) override;
nvinfer1::IPluginV2DynamicExt* deserializePlugin(
const char* name, const void* serialData, size_t serialLength) override;
void setPluginNamespace(const char* pluginNamespace) override;
const char* getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_INSTANCE_NORMALIZATION_PLUGIN_H
#ifndef TRT_MODULATED_DEFORM_CONV_HPP
#define TRT_MODULATED_DEFORM_CONV_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
class ModulatedDeformableConvPluginDynamic
: public nvinfer1::IPluginV2DynamicExt {
public:
ModulatedDeformableConvPluginDynamic(const std::string &name,
const nvinfer1::Dims stride,
const nvinfer1::Dims padding,
const nvinfer1::Dims dilation,
const int deformableGroup,
const int group);
ModulatedDeformableConvPluginDynamic(const std::string name, const void *data,
size_t length);
ModulatedDeformableConvPluginDynamic() = delete;
~ModulatedDeformableConvPluginDynamic();
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
void attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
nvinfer1::IGpuAllocator *gpuAllocator) override;
void detachFromContext() override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
nvinfer1::Dims mStride;
nvinfer1::Dims mPadding;
nvinfer1::Dims mDilation;
int mDeformableGroup;
int mGroup;
bool mWithBias;
cublasHandle_t m_cublas_handle;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class ModulatedDeformableConvPluginDynamicCreator
: public nvinfer1::IPluginCreator {
public:
ModulatedDeformableConvPluginDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_MODULATED_DEFORM_CONV_HPP
#ifndef TRT_NMS_HPP
#define TRT_NMS_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
class NonMaxSuppressionDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
NonMaxSuppressionDynamic(const std::string &name, int centerPointBox,
int maxOutputBoxesPerClass, float iouThreshold,
float scoreThreshold, int offset);
NonMaxSuppressionDynamic(const std::string name, const void *data,
size_t length);
NonMaxSuppressionDynamic() = delete;
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
int mCenterPointBox;
int mMaxOutputBoxesPerClass;
float mIouThreshold;
float mScoreThreshold;
int mOffset;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class NonMaxSuppressionDynamicCreator : public nvinfer1::IPluginCreator {
public:
NonMaxSuppressionDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_NMS_HPP
#ifndef TRT_PLUGIN_HPP
#define TRT_PLUGIN_HPP
extern "C" {
bool initLibMMCVInferPlugins();
} // extern "C"
#endif // TRT_PLUGIN_HPP
#ifndef TRT_PLUGIN_HELPER_HPP
#define TRT_PLUGIN_HELPER_HPP
#include <stdexcept>
#include "NvInferPlugin.h"
namespace mmcv {
const int MAXTENSORDIMS = 10;
struct TensorDesc {
int shape[MAXTENSORDIMS];
int stride[MAXTENSORDIMS];
int dim;
};
inline unsigned int getElementSize(nvinfer1::DataType t) {
switch (t) {
case nvinfer1::DataType::kINT32:
return 4;
case nvinfer1::DataType::kFLOAT:
return 4;
case nvinfer1::DataType::kHALF:
return 2;
// case nvinfer1::DataType::kBOOL:
case nvinfer1::DataType::kINT8:
return 1;
default:
throw std::runtime_error("Invalid DataType.");
}
throw std::runtime_error("Invalid DataType.");
return 0;
}
inline size_t getAlignedSize(size_t origin_size, size_t aligned_number = 16) {
return size_t((origin_size + aligned_number - 1) / aligned_number) *
aligned_number;
}
} // namespace mmcv
#endif // TRT_PLUGIN_HELPER_HPP
#ifndef TRT_ROI_ALIGN_HPP
#define TRT_ROI_ALIGN_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
class RoIAlignPluginDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
RoIAlignPluginDynamic(const std::string &name, int outWidth, int outHeight,
float spatialScale, int sampleRatio, int poolMode,
bool aligned);
RoIAlignPluginDynamic(const std::string name, const void *data,
size_t length);
RoIAlignPluginDynamic() = delete;
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
int mOutWidth;
int mOutHeight;
float mSpatialScale;
int mSampleRatio;
int mPoolMode; // 1:avg 0:max
bool mAligned;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class RoIAlignPluginDynamicCreator : public nvinfer1::IPluginCreator {
public:
RoIAlignPluginDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_ROI_ALIGN_HPP
#ifndef TRT_SCATTERND_HPP
#define TRT_SCATTERND_HPP
#include <cublas_v2.h>
#include <memory>
#include <string>
#include <vector>
#include "trt_plugin_helper.hpp"
class ONNXScatterNDDynamic : public nvinfer1::IPluginV2DynamicExt {
public:
ONNXScatterNDDynamic(const std::string &name);
ONNXScatterNDDynamic(const std::string name, const void *data, size_t length);
ONNXScatterNDDynamic() = delete;
// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt *clone() const override;
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs,
nvinfer1::IExprBuilder &exprBuilder) override;
bool supportsFormatCombination(int pos,
const nvinfer1::PluginTensorDesc *inOut,
int nbInputs, int nbOutputs) override;
void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
int nbInputs,
const nvinfer1::DynamicPluginTensorDesc *out,
int nbOutputs) override;
size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
int nbInputs,
const nvinfer1::PluginTensorDesc *outputs,
int nbOutputs) const override;
int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
const nvinfer1::PluginTensorDesc *outputDesc,
const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(int index,
const nvinfer1::DataType *inputTypes,
int nbInputs) const override;
// IPluginV2 Methods
const char *getPluginType() const override;
const char *getPluginVersion() const override;
int getNbOutputs() const override;
int initialize() override;
void terminate() override;
size_t getSerializationSize() const override;
void serialize(void *buffer) const override;
void destroy() override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
const std::string mLayerName;
std::string mNamespace;
protected:
// To prevent compiler warnings.
using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
};
class ONNXScatterNDDynamicCreator : public nvinfer1::IPluginCreator {
public:
ONNXScatterNDDynamicCreator();
const char *getPluginName() const override;
const char *getPluginVersion() const override;
const nvinfer1::PluginFieldCollection *getFieldNames() override;
nvinfer1::IPluginV2 *createPlugin(
const char *name, const nvinfer1::PluginFieldCollection *fc) override;
nvinfer1::IPluginV2 *deserializePlugin(const char *name,
const void *serialData,
size_t serialLength) override;
void setPluginNamespace(const char *pluginNamespace) override;
const char *getPluginNamespace() const override;
private:
static nvinfer1::PluginFieldCollection mFC;
static std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
#endif // TRT_SCATTERND_HPP
// Modified from:
// https://github.com/NVIDIA/TensorRT/blob/master/plugin/common/serialize.hpp
#ifndef TRT_SERIALIZE_HPP
#define TRT_SERIALIZE_HPP
#include <cassert>
#include <cstring>
#include <iostream>
#include <type_traits>
#include <vector>
using std::cerr;
using std::cout;
using std::endl;
template <typename T>
inline void serialize_value(void** buffer, T const& value);
template <typename T>
inline void deserialize_value(void const** buffer, size_t* buffer_size,
T* value);
namespace {
template <typename T, class Enable = void>
struct Serializer {};
template <typename T>
struct Serializer<T, typename std::enable_if<std::is_arithmetic<T>::value ||
std::is_enum<T>::value ||
std::is_pod<T>::value>::type> {
static size_t serialized_size(T const& value) { return sizeof(T); }
static void serialize(void** buffer, T const& value) {
::memcpy(*buffer, &value, sizeof(T));
reinterpret_cast<char*&>(*buffer) += sizeof(T);
}
static void deserialize(void const** buffer, size_t* buffer_size, T* value) {
assert(*buffer_size >= sizeof(T));
::memcpy(value, *buffer, sizeof(T));
reinterpret_cast<char const*&>(*buffer) += sizeof(T);
*buffer_size -= sizeof(T);
}
};
template <>
struct Serializer<const char*> {
static size_t serialized_size(const char* value) { return strlen(value) + 1; }
static void serialize(void** buffer, const char* value) {
::strcpy(static_cast<char*>(*buffer), value);
reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
}
static void deserialize(void const** buffer, size_t* buffer_size,
const char** value) {
*value = static_cast<char const*>(*buffer);
size_t data_size = strnlen(*value, *buffer_size) + 1;
assert(*buffer_size >= data_size);
reinterpret_cast<char const*&>(*buffer) += data_size;
*buffer_size -= data_size;
}
};
template <typename T>
struct Serializer<std::vector<T>,
typename std::enable_if<std::is_arithmetic<T>::value ||
std::is_enum<T>::value ||
std::is_pod<T>::value>::type> {
static size_t serialized_size(std::vector<T> const& value) {
return sizeof(value.size()) + value.size() * sizeof(T);
}
static void serialize(void** buffer, std::vector<T> const& value) {
serialize_value(buffer, value.size());
size_t nbyte = value.size() * sizeof(T);
::memcpy(*buffer, value.data(), nbyte);
reinterpret_cast<char*&>(*buffer) += nbyte;
}
static void deserialize(void const** buffer, size_t* buffer_size,
std::vector<T>* value) {
size_t size;
deserialize_value(buffer, buffer_size, &size);
value->resize(size);
size_t nbyte = value->size() * sizeof(T);
assert(*buffer_size >= nbyte);
::memcpy(value->data(), *buffer, nbyte);
reinterpret_cast<char const*&>(*buffer) += nbyte;
*buffer_size -= nbyte;
}
};
} // namespace
template <typename T>
inline size_t serialized_size(T const& value) {
return Serializer<T>::serialized_size(value);
}
template <typename T>
inline void serialize_value(void** buffer, T const& value) {
return Serializer<T>::serialize(buffer, value);
}
template <typename T>
inline void deserialize_value(void const** buffer, size_t* buffer_size,
T* value) {
return Serializer<T>::deserialize(buffer, buffer_size, value);
}
#endif // TRT_SERIALIZE_HPP
......@@ -17,16 +17,18 @@ class DeformRoIPoolFunction(Function):
@staticmethod
def symbolic(g, input, rois, offset, output_size, spatial_scale,
sampling_ratio, gamma):
inputs = [input, rois]
if offset is not None:
inputs = [input, rois, offset]
return g.op(
'mmcv::MMCVDeformRoIPool',
input,
rois,
offset,
*inputs,
pooled_height_i=output_size[0],
pooled_width_i=output_size[1],
spatial_scale_f=spatial_scale,
sampling_ratio_f=sampling_ratio,
gamma_f=gamma)
gamma_f=gamma,
)
@staticmethod
def forward(ctx,
......
......@@ -16,19 +16,6 @@ ext_module = ext_loader.load_ext('_ext', [
class SigmoidFocalLossFunction(Function):
@staticmethod
def symbolic(g, input: torch.Tensor, target: torch.LongTensor,
gamma: float, alpha: float, weight: torch.Tensor,
reduction: str):
return g.op(
'mmcv::MMCVSigmoidFocalLoss',
input,
target,
gamma_f=gamma,
alpha_f=alpha,
weight_f=weight,
reduction_s=reduction)
@staticmethod
def forward(ctx,
input: torch.Tensor,
......@@ -121,19 +108,6 @@ class SigmoidFocalLoss(nn.Module):
class SoftmaxFocalLossFunction(Function):
@staticmethod
def symbolic(g, input: torch.Tensor, target: torch.LongTensor,
gamma: float, alpha: float, weight: torch.Tensor,
reduction: str):
return g.op(
'mmcv::MMCVSoftmaxFocalLoss',
input,
target,
gamma_f=gamma,
alpha_f=alpha,
weight_f=weight,
reduction_s=reduction)
@staticmethod
def forward(ctx,
input: torch.Tensor,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment