def register_elementwise_kernel(
OperatorClass: Type[Elementwise],
code: str,
parameters: Dict[str, Callable[[Elementwise], Union[int,
float]]] = None):
"""
Utility function to define elementwise operation kernel in Fallback backend.
In code generation phase, registered `OperatorClass` operator is handled by
:func:`webdnn.backend.fallback.kernels.elementwise_kernel<elementwise_kernel>`, and kernel code is generated with `code`.::
# With expression code
register_elementwise_kernel(ElementwiseAdd, "y = x0 + x1;")
# With block code
register_elementwise_kernel(HardSigmoid, \"\"\"
y = x0 * 0.2 + 0.5;
if (y < 0.0) {
y = 0.0;
} else if (y > 1.0) {
y = 1.0;
}
\"\"\")
# With hyper parameters
register_elementwise_kernel(ClippedRelu,
"y = x0 < 0 ? 0 : x0 > cap ? cap : x0;",
{ "cap": lambda op: op.parameters["cap"] })
Like above examples, `code` is only partial code to compute single output element. Other procedures such as memory load and store
operation is not needed. variable `x0`, `x1`, ..., `x_{n-1}` and `y` is pre-defined, where `n` is number of input variables.
If you want to use hyper parameters, use `parameters` argument like last one of above examples. Only int and float value are supported.
Args:
OperatorClass: Operator class which the handler is bound to
code: Operator code in JavaScript
parameters: Hyper parameters
"""
FallbackDescriptorGenerator.register_handler(OperatorClass)(
elementwise_kernel)
_registered_items[OperatorClass] = RegisteredItem(
OperatorClass=OperatorClass,
code=code,
parameters={} if parameters is None else parameters)
import keras
import square
from webdnn.backend.fallback.generator import FallbackDescriptorGenerator
from webdnn.backend.webassembly.generator import WebassemblyDescriptorGenerator
from webdnn.backend.webgpu.generator import WebGPUDescriptorGenerator
from webdnn.frontend.keras.converter import KerasConverter
# ---------------------------------------------------------------------------------------------------
# Define Keras model
x = keras.layers.Input((10, ))
y = square.SquareLayer()(x)
model = keras.models.Model([x], [y])
# ---------------------------------------------------------------------------------------------------
# Convert Keras model into WebDNN graph IR
graph = KerasConverter(batch_size=1).convert(model)
# ---------------------------------------------------------------------------------------------------
# Generate graph descriptors
WebGPUDescriptorGenerator.generate(graph).save("./output")
WebassemblyDescriptorGenerator.generate(graph).save("./output")
FallbackDescriptorGenerator.generate(graph).save("./output")