Custom deep learning layers support

Original author	Dmitry Kurtaev
Compatibility	OpenCV >= 3.4.1

Introduction

Deep learning is a fast-growing area. New approaches to building neural networks usually introduce new types of layers. These could be modifications of existing ones or implementation of outstanding research ideas.

OpenCV allows importing and running networks from different deep learning frameworks. There is a number of the most popular layers. However, you can face a problem that your network cannot be imported using OpenCV because some layers of your network can be not implemented in the deep learning engine of OpenCV.

The first solution is to create a feature request at opencv/opencv#issues mentioning details such as a source of a model and a type of new layer. The new layer could be implemented if the OpenCV community shares this need.

The second way is to define a custom layer so that OpenCV’s deep learning engine will know how to use it. This tutorial is dedicated to show you a process of deep learning model’s import customization.

Define a custom layer in C++

Deep learning layer is a building block of network’s pipeline. It has connections to input blobs and produces results to output blobs. There are trained weights and hyper-parameters. Layers’ names, types, weights and hyper-parameters are stored in files are generated by native frameworks during training. If OpenCV encounters unknown layer type it throws an exception while trying to read a model:

Unspecified error: Can't create layer "layer_name" of type "MyType" in function getLayerInstance

To import the model correctly you have to derive a class from cv::dnn::Layer with the following methods:

class MyLayer : public cv::dnn::Layer
{
public:
    //! [MyLayer::MyLayer]
    MyLayer(const cv::dnn::LayerParams &params);
    //! [MyLayer::MyLayer]

    //! [MyLayer::create]
    static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params);
    //! [MyLayer::create]

    //! [MyLayer::getMemoryShapes]
    virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
                                 const int requiredOutputs,
                                 std::vector<std::vector<int> > &outputs,
                                 std::vector<std::vector<int> > &internals) const CV_OVERRIDE;
    //! [MyLayer::getMemoryShapes]

    //! [MyLayer::forward]
    virtual void forward(cv::InputArrayOfArrays inputs,
                         cv::OutputArrayOfArrays outputs,
                         cv::OutputArrayOfArrays internals) CV_OVERRIDE;
    //! [MyLayer::forward]

    //! [MyLayer::finalize]
    virtual void finalize(cv::InputArrayOfArrays inputs,
                          cv::OutputArrayOfArrays outputs) CV_OVERRIDE;
    //! [MyLayer::finalize]
};

And register it before the import:

#include <opencv2/dnn/layer.details.hpp>  // CV_DNN_REGISTER_LAYER_CLASS

static inline void loadNet()
{
    CV_DNN_REGISTER_LAYER_CLASS(Interp, InterpLayer);
    // ...

Note

MyType is a type of unimplemented layer from the thrown exception.

Let’s see what all the methods do:

• Constructor

MyLayer(const cv::dnn::LayerParams &params);

Retrieves hyper-parameters from cv::dnn::LayerParams. If your layer has trainable weights they will be already stored in the Layer’s member cv::dnn::Layer::blobs.

• A static method create

static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params);

This method should create an instance of you layer and return cv::Ptr with it.

• Output blobs’ shape computation

virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
                             const int requiredOutputs,
                             std::vector<std::vector<int> > &outputs,
                             std::vector<std::vector<int> > &internals) const CV_OVERRIDE;

Returns layer’s output shapes depending on input shapes. You may request an extra memory using internals.

• Run a layer

virtual void forward(cv::InputArrayOfArrays inputs,
                     cv::OutputArrayOfArrays outputs,
                     cv::OutputArrayOfArrays internals) CV_OVERRIDE;

Implement a layer’s logic here. Compute outputs for given inputs.

Note

OpenCV manages memory allocated for layers. In the most cases the same memory can be reused between layers. So your forward implementation should not rely on that the second invocation of forward will have the same data at outputs and internals.

• Optional finalize method

virtual void finalize(cv::InputArrayOfArrays inputs,
                      cv::OutputArrayOfArrays outputs) CV_OVERRIDE;

The chain of methods is the following: OpenCV deep learning engine calls create method once, then it calls getMemoryShapes for every created layer, then you can make some preparations depend on known input dimensions at cv::dnn::Layer::finalize. After network was initialized only forward method is called for every network’s input.

Note

Varying input blobs’ sizes such height, width or batch size make OpenCV reallocate all the internal memory. That leads to efficiency gaps. Try to initialize and deploy models using a fixed batch size and image’s dimensions.

Example: custom layer from TensorFlow

This is an example of how to import a network with tf.image.resize_bilinear operation. This is also a resize but with an implementation different from OpenCV’s built-in resize.

Let’s create a single layer network:

inp = tf.placeholder(tf.float32, [2, 3, 4, 5], 'input')
resized = tf.image.resize_bilinear(inp, size=[9, 8], name='resize_bilinear')

OpenCV sees that TensorFlow’s graph in the following way:

node {
  name: "input"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
}
node {
  name: "resize_bilinear/size"
  op: "Const"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_INT32
        tensor_shape {
          dim {
            size: 2
          }
        }
        tensor_content: "\t\000\000\000\010\000\000\000"
      }
    }
  }
}
node {
  name: "resize_bilinear"
  op: "ResizeBilinear"
  input: "input:0"
  input: "resize_bilinear/size"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "align_corners"
    value {
      b: false
    }
  }
}
library {
}

Custom layers import from TensorFlow is designed to put all layer’s attr into cv::dnn::LayerParams but input Const blobs into cv::dnn::Layer::blobs. In our case resize’s output shape will be stored in layer’s blobs[0].

class ResizeBilinearLayer CV_FINAL : public cv::dnn::Layer
{
public:
    ResizeBilinearLayer(const cv::dnn::LayerParams &params) : Layer(params)
    {
        CV_Assert(!params.get<bool>("align_corners", false));
        CV_Assert(!blobs.empty());

        for (size_t i = 0; i < blobs.size(); ++i)
            CV_Assert(blobs[i].type() == CV_32SC1);

        // There are two cases of input blob: a single blob which contains output
        // shape and two blobs with scaling factors.
        if (blobs.size() == 1)
        {
            CV_Assert(blobs[0].total() == 2);
            outHeight = blobs[0].at<int>(0, 0);
            outWidth = blobs[0].at<int>(0, 1);
            factorHeight = factorWidth = 0;
        }
        else
        {
            CV_Assert(blobs.size() == 2); CV_Assert(blobs[0].total() == 1); CV_Assert(blobs[1].total() == 1);
            factorHeight = blobs[0].at<int>(0, 0);
            factorWidth = blobs[1].at<int>(0, 0);
            outHeight = outWidth = 0;
        }
    }

    static cv::Ptr<cv::dnn::Layer> create(cv::dnn::LayerParams& params)
    {
        return cv::Ptr<cv::dnn::Layer>(new ResizeBilinearLayer(params));
    }

    virtual bool getMemoryShapes(const std::vector<std::vector<int> > &inputs,
                                 const int,
                                 std::vector<std::vector<int> > &outputs,
                                 std::vector<std::vector<int> > &) const CV_OVERRIDE
    {
        std::vector<int> outShape(4);
        outShape[0] = inputs[0][0];  // batch size
        outShape[1] = inputs[0][1];  // number of channels
        outShape[2] = outHeight != 0 ? outHeight : (inputs[0][2] * factorHeight);
        outShape[3] = outWidth != 0 ? outWidth : (inputs[0][3] * factorWidth);
        outputs.assign(1, outShape);
        return false;
    }

    virtual void finalize(cv::InputArrayOfArrays, cv::OutputArrayOfArrays outputs_arr) CV_OVERRIDE
    {
        std::vector<cv::Mat> outputs;
        outputs_arr.getMatVector(outputs);
        if (!outWidth && !outHeight)
        {
            outHeight = outputs[0].size[2];
            outWidth = outputs[0].size[3];
        }
    }

    // This implementation is based on a reference implementation from
    // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h
    virtual void forward(cv::InputArrayOfArrays inputs_arr,
                         cv::OutputArrayOfArrays outputs_arr,
                         cv::OutputArrayOfArrays internals_arr) CV_OVERRIDE
    {
        if (inputs_arr.depth() == CV_16S)
        {
            // In case of DNN_TARGET_OPENCL_FP16 target the following method
            // converts data from FP16 to FP32 and calls this forward again.
            forward_fallback(inputs_arr, outputs_arr, internals_arr);
            return;
        }

        std::vector<cv::Mat> inputs, outputs;
        inputs_arr.getMatVector(inputs);
        outputs_arr.getMatVector(outputs);

        cv::Mat& inp = inputs[0];
        cv::Mat& out = outputs[0];
        const float* inpData = (float*)inp.data;
        float* outData = (float*)out.data;

        const int batchSize = inp.size[0];
        const int numChannels = inp.size[1];
        const int inpHeight = inp.size[2];
        const int inpWidth = inp.size[3];

        float heightScale = static_cast<float>(inpHeight) / outHeight;
        float widthScale = static_cast<float>(inpWidth) / outWidth;
        for (int b = 0; b < batchSize; ++b)
        {
            for (int y = 0; y < outHeight; ++y)
            {
                float input_y = y * heightScale;
                int y0 = static_cast<int>(std::floor(input_y));
                int y1 = std::min(y0 + 1, inpHeight - 1);
                for (int x = 0; x < outWidth; ++x)
                {
                    float input_x = x * widthScale;
                    int x0 = static_cast<int>(std::floor(input_x));
                    int x1 = std::min(x0 + 1, inpWidth - 1);
                    for (int c = 0; c < numChannels; ++c)
                    {
                        float interpolation =
                            inpData[offset(inp.size, c, x0, y0, b)] * (1 - (input_y - y0)) * (1 - (input_x - x0)) +
                            inpData[offset(inp.size, c, x0, y1, b)] * (input_y - y0) * (1 - (input_x - x0)) +
                            inpData[offset(inp.size, c, x1, y0, b)] * (1 - (input_y - y0)) * (input_x - x0) +
                            inpData[offset(inp.size, c, x1, y1, b)] * (input_y - y0) * (input_x - x0);
                        outData[offset(out.size, c, x, y, b)] = interpolation;
                    }
                }
            }
        }
    }

private:
    static inline int offset(const cv::MatSize& size, int c, int x, int y, int b)
    {
        return x + size[3] * (y + size[2] * (c + size[1] * b));
    }

    int outWidth, outHeight, factorWidth, factorHeight;
};

Next we register a layer and try to import the model.

CV_DNN_REGISTER_LAYER_CLASS(ResizeBilinear, ResizeBilinearLayer);
cv::dnn::Net tfNet = cv::dnn::readNet("/path/to/graph.pb");

Example: custom layer from ONNX

ONNX groups operators into domains. The standard operators live in the default domain ai.onnx; vendors and exporters often place their own ops in a named domain such as my.namespace. When OpenCV imports an ONNX node, it looks the op up in cv::dnn::LayerFactory by:

• the op_type alone, for nodes in the default ai.onnx domain (or no domain), and

• "<domain>.<op_type>", for nodes in any non-default domain.

Node attributes are passed through to the layer constructor as cv::dnn::LayerParams entries with the same names. Consider an op MyCustomOp with attributes scale and bias that computes y = scale * x + bias. The implementation can look like:

// y = scale * x + bias, with scale/bias read from ONNX node attributes.
class CustomScaleBiasLayer CV_FINAL : public Layer
{
public:
    CustomScaleBiasLayer(const LayerParams& params) : Layer(params)
    {
        scale = params.get<float>("scale", 1.f);
        bias  = params.get<float>("bias",  0.f);
    }

    static Ptr<Layer> create(LayerParams& params)
    {
        return makePtr<CustomScaleBiasLayer>(params);
    }

    bool getMemoryShapes(const vector<MatShape>& inpts,
                         const int /*requiredOutputs*/,
                         vector<MatShape>& outShapes,
                         vector<MatShape>& /*internals*/) const CV_OVERRIDE
    {
        outShapes.assign(1, inpts[0]);
        return false;
    }

    void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr,
                 OutputArrayOfArrays) CV_OVERRIDE
    {
        vector<Mat> inps, outs;
        inputs_arr.getMatVector(inps);
        outputs_arr.getMatVector(outs);
        inps[0].convertTo(outs[0], outs[0].type(), scale, bias);
    }

private:
    float scale, bias;
};

To import a model that uses this op, register the layer before calling cv::dnn::readNetFromONNX. Use cv::dnn::LayerFactory::registerLayer for runtime registration (and cv::dnn::LayerFactory::unregisterLayer when done) — pick the right key for the domain of the op as described above:

// ONNX op-type lookup: layers in the default `ai.onnx` domain are registered
// under their op_type; layers in a non-default domain are registered under
// "<domain>.<op_type>" (e.g. "my.namespace.MyDomainOp").
LayerFactory::registerLayer(opKey, CustomScaleBiasLayer::create);

A complete runnable example is available at samples/dnn/custom_layer_onnx.cpp. Tiny ONNX models exercising both the default-domain and custom-domain registration paths can be generated with generate_custom_layer_models.py in the opencv_extra repository.

Define a custom layer in Python

The following example shows how to customize OpenCV’s layers in Python.

Let’s consider the Holistically-Nested Edge Detection model. Its Crop layers receive two input blobs and crop the first one to match the spatial dimensions of the second. OpenCV’s built-in Crop layer trims from the top-left corner, whereas this model expects cropping from the center, so using the built-in behaviour directly would produce shifted results with filled borders.

Next we’re going to replace OpenCV’s Crop layer that makes top-left cropping by a centric one.

• Create a class with getMemoryShapes and forward methods

class CropLayer(object):
    def __init__(self, params, blobs):
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    # Our layer receives two inputs. We need to crop the first input blob
    # to match a shape of the second one (keeping batch size and number of channels)
    def getMemoryShapes(self, inputs):
        inputShape, targetShape = inputs[0], inputs[1]
        batchSize, numChannels = inputShape[0], inputShape[1]
        height, width = targetShape[2], targetShape[3]

        self.ystart = (inputShape[2] - targetShape[2]) // 2
        self.xstart = (inputShape[3] - targetShape[3]) // 2
        self.yend = self.ystart + height
        self.xend = self.xstart + width

        return [[batchSize, numChannels, height, width]]

    def forward(self, inputs):
        return [inputs[0][:,:,self.ystart:self.yend,self.xstart:self.xend]]

Note

Both methods should return lists.

• Register a new layer.

cv.dnn_registerLayer('Crop', CropLayer)

That’s it! We have replaced an implemented OpenCV’s layer to a custom one. You may find a full script in the source code.