Support for Custom Layers Defined in Python

[Willaaaaaaa]

Summary

Our current implementation of Net.register_custom_layer method is not usable and fully-featured. This api's goal is to allow users to define custom neural network layers in Python (by subclassing ncnn_mp.Layer) and register them with an ncnn network.

Current Status

ncnn_mp/modules/stub/ncnn_mp.pyi

Lines 676 to 688 in 3584975

	@overload
	def register_custom_layer(self, type: str, layer_class: Type[Layer]) -> None: ...
	@overload
	def register_custom_layer(self, typeindex: int, layer_class: Type[Layer]) -> None: ...
	def register_custom_layer(self, identifier, layer_class):
	"""
	DO NOT USE THIS API
	Register a custom layer implementation.
	**WARNING**: This API is WIP and unusable from Python.
	"""
	...

• Pythonic API design is complete: The ideal Python interface and usage pattern have been designed (see below).
• C-level implementation is currently paused: The underlying code in ncnn_mp.c that bridges MicroPython and ncnn is not yet implemented because of the complexity between the MicroPython C API and the ncnn design.

The Challenge

The core challenge is bridging the ncnn C API's function-pointer callback mechanism with MicroPython.

The final implementation should works like this:

Register:
User python statement writing -> call api binding in ncnn_mp.c -> run the trampoline function in ncnn_mp.c to process the method and property in MyLayer class -> call the underlying ncnn implementation

And in excution period, when Python run extractor.extract(), ncnn can invoke extractor to run the layer.

Desired Python API

The goal is to enable the following workflow for users:

import ncnn_mp

# A user-defined custom layer
class MyLayer(ncnn_mp.Layer):
    def __init__(self):
        super().__init__()
        self.one_blob_only = True
        self.support_inplace = True
    
    def forward_inplace(self, bottom_top_blob, opt):
        # Your logic...
        return 0

    # And other methods users need

# User registers the custom class with a Net instance
my_net = ncnn_mp.Net()
my_net.register_custom_layer("MyLayer", MyLayer)

# ...

Next Steps

• Carefully read the ncnn source code, specifically Net::register_custom_layer and c_api.cpp, to fully understand the mechanics behind.
• Implement this api.