Extending Model Optimizer with Caffe* Python Layers

This section provides instruction on how to support a custom Caffe operation written only in Python. For example, the Faster-R-CNN model implemented in Caffe contains a custom layer Proposal written in Python. The layer is described in the Faster-R-CNN protoxt the following way:

layer {
name: 'proposal'
type: 'Python'
bottom: 'rpn_cls_prob_reshape'
bottom: 'rpn_bbox_pred'
bottom: 'im_info'
top: 'rois'
python_param {
module: 'rpn.proposal_layer'
layer: 'ProposalLayer'
param_str: "'feat_stride': 16"

This section describes only a procedure on how to extract operator attributes in the Model Optimizer. The rest of the operation enabling pipeline and documentation on how to support other Caffe operations (written in C++) is described in the main document Customize_Model_Optimizer.

Writing Extractor for Caffe Python Layer

Custom Caffe Python layers have an attribute type (defining the type of the operation) equal to Python and two mandatory attributes module and layer in the python_param dictionary. The module defines the Python module name with the layer implementation, while layer value is an operation type defined by an user. In order to extract attributes for such an operation it is necessary to implement extractor class inherited from the CaffePythonFrontExtractorOp class instead of FrontExtractorOp class used for standard framework layers. The op class attribute value should be set to the module + "." + layer value so the extractor is triggered for this kind of operation.

Here is a simplified example of the extractor for the custom operation Proposal from Faster-R-CNN model mentioned above. The full code with additional checks is provided in the <INSTALL_DIR>/deployment_tools/model_optimizer/extensions/front/caffe/proposal_python_ext.py. The sample code uses operation ProposalOp which corresponds to Proposal operation described in the Available Operations Sets document. Refer to the source code below for a detailed explanation of the extractor.

from extensions.ops.proposal import ProposalOp
from mo.front.extractor import CaffePythonFrontExtractorOp
class ProposalPythonFrontExtractor(CaffePythonFrontExtractorOp):
op = 'rpn.proposal_layer.ProposalLayer' # module + "." + layer
enabled = True # extractor is enabled
def extract_proposal_params(node, defaults):
param = node.pb.python_param # get the protobuf message representation of the layer attributes
# parse attributes from the layer protobuf message to a Python dictionary
attrs = CaffePythonFrontExtractorOp.parse_param_str(param.param_str)
update_attrs = defaults
# the operation expects ratio and scale values to be called "ratio" and "scale" while Caffe uses different names
if 'ratios' in attrs:
attrs['ratio'] = attrs['ratios']
del attrs['ratios']
if 'scales' in attrs:
attrs['scale'] = attrs['scales']
del attrs['scales']
ProposalOp.update_node_stat(node, update_attrs) # update the node attributes
def extract(cls, node):
# define default values for the Proposal layer attributes
defaults = {
'feat_stride': 16,
'base_size': 16,
'min_size': 16,
'ratio': [0.5, 1, 2],
'scale': [8, 16, 32],
'pre_nms_topn': 6000,
'post_nms_topn': 300,
'nms_thresh': 0.7
cls.extract_proposal_params(node, defaults)
return cls.enabled

See Also