I implemented a custom layer in c++ and cuda by inheriting from PluginV2, and creating a PluginCreator. Since I generate my engine in python, I used pybind11 to bind the ctors of the layer and the creator. In python, I use network.add_plugin_v2() to add the layers to my network.
The problem is that every time the engine calls the enqueue() function of my layer, it always calls a single instance of the layer when there are about 10 different ones… The reason it’s not working is because every instance of the layer is initialized with different values, so they are supposed to act differently.
this is the code for the plugin, its for a poc so ignore the prints and lack of standards
#include "SplitPlugin.h"
#include "NvInfer.h"
#include "splitKernel.h"
#include <vector>
#include <cassert>
#include <cstring>
#include <iostream>
using namespace nvinfer1;
// Split plugin specific constants
namespace {
static const char* SPLIT_PLUGIN_VERSION{"1"};
static const char* SPLIT_PLUGIN_NAME{"SplitPlugin"};
}
// Static class fields initialization
PluginFieldCollection SplitPluginCreator::mFC{};
std::vector<PluginField> SplitPluginCreator::mPluginAttributes;
REGISTER_TENSORRT_PLUGIN(SplitPluginCreator);
// Helper function for serializing plugin
template<typename T>
void writeToBuffer(char*& buffer, const T& val)
{
*reinterpret_cast<T*>(buffer) = val;
buffer += sizeof(T);
}
// Helper function for deserializing plugin
template<typename T>
T readFromBuffer(const char*& buffer)
{
T val = *reinterpret_cast<const T*>(buffer);
buffer += sizeof(T);
return val;
}
SplitPlugin::SplitPlugin(const std::string name)
: mLayerName(name), mInputVolume(0)
{
std::cout << "SplitPlugin called in " << this << std::endl;
}
SplitPlugin::SplitPlugin(const std::string name, const void* data, size_t length)
: mLayerName(name)
{
std::cout << "SplitPluginDeserialize called in " << this << std::endl;
// Deserialize in the same order as serialization
const char *d = static_cast<const char *>(data);
const char *a = d;
mInputVolume = readFromBuffer<size_t>(d);
assert(d == (a + length));
}
const char* SplitPlugin::getPluginType() const
{
std::cout << "getPluginType called in " << this << std::endl;
return SPLIT_PLUGIN_NAME;
}
const char* SplitPlugin::getPluginVersion() const
{
std::cout << "getPluginVersion called in " << this << std::endl;
return SPLIT_PLUGIN_VERSION;
}
int SplitPlugin::getNbOutputs() const
{
std::cout << "getNbOutputs called in " << this << std::endl;
return 2;
}
Dims SplitPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
{
std::cout << "getOutputDimensions called in " << this << std::endl;
// Validate input arguments
assert(nbInputDims == 1);
assert(inputs[0].nbDims == 3);
assert(index <= 1 && index >= 0);
Dims3 dims(inputs[0].d[0] / 2, inputs[0].d[1], inputs[0].d[2]);
// Assuming nchw, we cut the channels into two.
return dims;
}
int SplitPlugin::initialize()
{
std::cout << "initialize called in " << this << std::endl;
return 0;
}
int SplitPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void*, cudaStream_t stream)
{
std::cout << "enqueue called in " << this << std::endl;
// TODO: split
//status = clipInference(stream, mInputVolume * batchSize, mSplitMin, mSplitMax, inputs[0], output);
std::cout << mInputVolume << std::endl;
std::cout << "this enqueue: " << this << std::endl;
StripInference(stream, batchSize, inputs[0], outputs, mInputVolume);
return 0;
}
size_t SplitPlugin::getSerializationSize() const
{
std::cout << "getSerializationSize called in " << this << std::endl;
return sizeof(size_t);
}
void SplitPlugin::serialize(void* buffer) const
{
std::cout << "serialize called in " << this << std::endl;
char *d = static_cast<char *>(buffer);
const char *a = d;
writeToBuffer(d, mInputVolume);
assert(d == a + getSerializationSize());
}
void SplitPlugin::configureWithFormat(const Dims* inputs, int nbInputs, const Dims* outputs, int nbOutputs, DataType type, PluginFormat format, int)
{
std::cout << "configureWithFormat called in " << this << std::endl;
// Validate input arguments
assert(nbInputs == 1);
assert(nbOutputs == 2);
assert(format == PluginFormat::kNCHW);
// Fetch volume for future enqueue() operations
size_t volume = 1;
for (int i = 0; i < inputs->nbDims; i++) {
volume *= inputs->d[i];
}
mInputVolume = volume;
std::cout << "configureWithFormat Called with volume: " << volume << "on: " << this<<std::endl;
}
bool SplitPlugin::supportsFormat(DataType type, PluginFormat format) const
{
std::cout << "supportsFormat called in " << this << std::endl;
// This plugin only supports ordinary floats, and NCHW input format
if (type == DataType::kFLOAT && format == PluginFormat::kNCHW)
return true;
else
return false;
}
void SplitPlugin::terminate() {std::cout << "terminate called in " << this << std::endl;}
void SplitPlugin::destroy() {
std::cout << "destroy called in " << this << std::endl;
// This gets called when the network containing plugin is destroyed
delete this;
}
IPluginV2* SplitPlugin::clone() const
{
IPluginV2* tmpptr = new SplitPlugin(mLayerName);
((SplitPlugin*)tmpptr)->mInputVolume = this->mInputVolume;
std::cout << "clone called in " << this <<" to" << tmpptr << std::endl;
return tmpptr;
}
void SplitPlugin::setPluginNamespace(const char* libNamespace)
{
std::cout << "setPluginNamespace called in " << this << std::endl;
mNamespace = libNamespace;
}
const char* SplitPlugin::getPluginNamespace() const
{
std::cout << "getPluginNamespace called in " << this << std::endl;
return mNamespace.c_str();
}
SplitPluginCreator::SplitPluginCreator()
{
// Describe SplitPlugin's required PluginField arguments
// Fill PluginFieldCollection with PluginField arguments metadata
mFC.nbFields = mPluginAttributes.size(); // 0.
mFC.fields = mPluginAttributes.data();
}
const char* SplitPluginCreator::getPluginName() const
{
return SPLIT_PLUGIN_NAME;
}
const char* SplitPluginCreator::getPluginVersion() const
{
return SPLIT_PLUGIN_VERSION;
}
const PluginFieldCollection* SplitPluginCreator::getFieldNames()
{
return &mFC;
}
IPluginV2* SplitPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
{
const PluginField* fields = fc->fields;
// Parse fields from PluginFieldCollection
assert(fc->nbFields == 0);
return new SplitPlugin(name);
}
IPluginV2* SplitPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
{
// This object will be deleted when the network is destroyed, which will
// call SplitPlugin::destroy()
return new SplitPlugin(name, serialData, serialLength);
}
void SplitPluginCreator::setPluginNamespace(const char* libNamespace)
{
mNamespace = libNamespace;
}
const char* SplitPluginCreator::getPluginNamespace() const
{
return mNamespace.c_str();
}
#include <pybind11/pybind11.h>
PYBIND11_MODULE(splitplugin, m)
{
namespace py = pybind11;
// This allows us to use the bindings exposed by the tensorrt module.
py::module::import("tensorrt");
// Note that we only need to bind the constructors manually. Since all other methods override IPlugin functionality, they will be automatically available in the python bindings.
// The `std::unique_ptr<SplitPlugin, py::nodelete>` specifies that Python is not responsible for destroying the object. This is required because the destructor is private.
py::class_<SplitPlugin, nvinfer1::IPluginV2, std::unique_ptr<SplitPlugin, py::nodelete>>(m, "SplitPlugin")
// Bind the normal constructor as well as the one which deserializes the plugin
.def(py::init<const std::string>())
.def(py::init<const std::string, const void*, size_t>())
;
py::class_<SplitPluginCreator, nvinfer1::IPluginCreator>(m, "SplitPluginCreator")
// Bind the default constructor.
.def(py::init<>())
;
}