Search code examples
caffematcaffe

Is it possible to add new layers to Caffe?


Is it possible for me to add a new layer (for a different loss) with Caffe? How can I do it using Matlab or Python?


Solution

  • Yes, you can add a custom loss function with pycaffe. Here is an example of Euclidean loss layer in python (taken from Caffe Github repo). You need to provide the loss function in forward function and it's gradient in backward method:

    import caffe
    import numpy as np
    
    class EuclideanLossLayer(caffe.Layer):
        """
        Compute the Euclidean Loss in the same manner as the C++ EuclideanLossLayer
        to demonstrate the class interface for developing layers in Python.
        """
    
        def setup(self, bottom, top):
            # check input pair
            if len(bottom) != 2:
                raise Exception("Need two inputs to compute distance.")
    
        def reshape(self, bottom, top):
            # check input dimensions match
            if bottom[0].count != bottom[1].count:
                raise Exception("Inputs must have the same dimension.")
            # difference is shape of inputs
            self.diff = np.zeros_like(bottom[0].data, dtype=np.float32)
            # loss output is scalar
            top[0].reshape(1)
    
        def forward(self, bottom, top):
            self.diff[...] = bottom[0].data - bottom[1].data
            top[0].data[...] = np.sum(self.diff**2) / bottom[0].num / 2.
    
        def backward(self, top, propagate_down, bottom):
            for i in range(2):
                if not propagate_down[i]:
                    continue
                if i == 0:
                    sign = 1
                else:
                    sign = -1
                bottom[i].diff[...] = sign * self.diff / bottom[i].num
    

    Save it as pyloss.py for example. You can then use the python layer in your prototxt file to load it:

    layer {
      type: 'Python'
      name: 'loss'
      top: 'loss'
      bottom: 'ipx'
      bottom: 'ipy'
      python_param {
        # the module name -- usually the filename -- that needs to be in $PYTHONPATH
        module: 'pyloss'
        # the layer name -- the class name in the module
        layer: 'EuclideanLossLayer'
      }
      # set loss weight so Caffe knows this is a loss layer.
      # since PythonLayer inherits directly from Layer, this isn't automatically
      # known to Caffe
      loss_weight: 1
    }
    

    Or in your python script:

    n.loss = L.Python(n.ipx, n.ipy,python_param=dict(module='pyloss',layer='EuclideanLossLayer'),
                         loss_weight=1)
    

    Just be very careful in computing and implemeting the gradient (backward function) since it tends to be error-prone.