TensorRT custom layer implementation: input dimension is zero (how to pass params?)

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

I’m trying to implement the Caffe Axpy layer using custom layer plugin, based on sample codes such as samplePlugin.cpp and Face Recognition. While checking every step for debugging, however, I found the input parameters are not transported correctly. Actually, I’m not sure how to implement the constructor part. Here is a little part of my code:

class Axpy : public nvinfer::IPluginExt
{
public:
   Axpy(){}
   Axpy(const void* data, size_t size)
   {
      const int* d = static_cast<const int*> (data), *a = d;
      dimsData = nvinfer1::DimsCHW{d[0], d[1], d[2]};
   }

   int getNbOutputs() const override
   {
      return 1;
   }

   Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) override
   {
      assert(mNbInputChannels == inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2]);
      return Dims3(mNbOutputChannels, 1, 1);
   }
   ...
   ...
protected:
   nvinfer1::DimsCHW dimsData;
};

// integration for serialization
class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactoryExt
{
public:
    // caffe parser plugin implementation
    bool isPlugin(const char *layerName) override
    {
        return isPluginExt(layerName);
    }

    bool isPluginExt(const char *layerName) override
    {
        return !strcmp(layerName, "conv2_1");
    }

    virtual nvinfer1::IPlugin* createPlugin(const char *layerName, const void *serialData, size_t serialLength) override
    {
        // there's no way to pass parameters through from the model definition, so we
        // have to define it here explicitly
        // TODO: do some relevant param pass, if any
        //const caffe::LayerParameter lm;
        mPlugin = std::unique_ptr<Axpy>(new Axpy(serialData, serialLength));
        return mPlugin.get();
    }

    // deserialization plugin implementation
    IPlugin* createPlugin(const char *layerName, const nvinfer1::Weights *weights, int nbWeights) override
    {
        assert(isPlugin(layerName));

        // this plugin object is destroyed when engine is destroyed by calling
        // IPluginExt::destroy()
        assert(mPlugin.get() == nullptr);
        mPlugin = std::unique_ptr<Axpy>(new Axpy());
        return mPlugin.get();
    }

    // user application destroys plugin when it is safe to do so.
    // should be done after consumers of plugin (like ICudaEngine) are destroyed.
    void destroyPlugin() { mPlugin.reset(); }
    std::unique_ptr<Axpy> mPlugin{ nullptr };

private:
    const nvinfer1::Weights* w;
    int nbWeights = 2;
    int nbOutputChannels = 3;
};

However, this gives me assertion failure for the line

assert(mNbInputChannels == inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2]);

After checking the values, I found

inputs[0].d[0] = 256; inputs[0].d[1] = 0; inputs[0].d[2] = 0;

I doubt I did not build the constructor correctly and buffers are not storing values accordingly. Interestingly, the code always enters the createPlugin of form:

IPlugin* createPlugin(const char *layerName, const nvinfer1::Weights *weights, int nbWeights)

but never enters

virtual nvinfer1::IPlugin* createPlugin(const char *layerName, const void *serialData, size_t serialLength) override

What did I miss here? And how to pass those values correctly? Thanks.