def test_python_extractor_for_extractors(self):
module = 'test_module'
layer = 'test_layer'
CaffePythonFrontExtractorOp.registered_ops['{}.{}'.format(module, layer)] = FakePythonExtractor
params = FakeMultiParam({
'module': module,
'layer': layer,
'param_str': "'feat_stride': 16"
})
self.assertTrue(PythonFrontExtractorOp.extract(FakeNode(FakePythonProtoLayer(params), None)))
def test_python_extractor_for_op(self):
module = 'test_module'
layer = 'test_layer'
CaffePythonFrontExtractorOp.registered_ops['{}.{}'.format(module, layer)] = \
lambda node: CaffePythonFrontExtractorOp.parse_param_str(node.pb.python_param.param_str)
params = FakeMultiParam({
'module': module,
'layer': layer,
'param_str': "'feat_stride': 16"
})
ext = PythonFrontExtractorOp.extract(FakeNode(FakePythonProtoLayer(params), None))
self.assertEqual({'feat_stride': 16}, ext)