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)