Models
HeartKit provides a number of model architectures that can be used for training heart-monitoring tasks. While a number of off-the-shelf models exist, they are often not efficient nor optimized for real-time, edge applications. To address this, HeartKit provides a model factory that allows you to easily create and train customized models via neuralspot-edge. neuralspot-edge includes a growing number of state-of-the-art models that can be easily configured and trained using high-level parameters. The models are designed to be efficient and well-suited for real-time, edge applications. Most of the models are based on state-of-the-art architectures that have been modified to allow for more fine-grain customization. In addition, the models support 1D variants to allow for training on time-series data. Please check neuralspot-edge for list of available models and their configurations.
Available Models
- TCN: A CNN leveraging dilated convolutions (key=
tcn) - U-Net: A CNN with encoder-decoder architecture for segmentation tasks (key=
unet) - U-NeXt: A U-Net variant leveraging MBConv blocks (key=
unext) - EfficientNetV2: A CNN leveraging MBConv blocks (key=
efficientnet) - MobileOne: A CNN aimed at sub-1ms inference (key=
mobileone) - ResNet: A popular CNN often used for vision tasks (key=
resnet) - Conformer: A transformer composed of both convolutional and self-attention blocks (key=
conformer) - MetaFormer: A transformer composed of both spatial mixing and channel mixing blocks (key=
metaformer) - TSMixer: An All-MLP Architecture for Time Series Classification (key=
tsmixer) - Bring-Your-Own-Model: Add a custom model architecture to HeartKit.
Model Factory
HeartKit includes a model factory, ModelFactory, that eases the processes of creating models for training. The factory allows you to create models by specifying the model key and the model parameters. The factory will then create the model using the specified parameters. The factory also allows you to register custom models that can be used for training. By leveraring a factory, a task only needs to provide the architecture key and the parameters, and the factory will take care of the rest.
The model factory provides the following methods:
- hk.ModelFactory.register: Register a custom model
- hk.ModelFactory.unregister: Unregister a custom model
- hk.ModelFactory.has: Check if a model is registered
- hk.ModelFactory.get: Get a model from the factory
- hk.ModelFactory.list: List all available models
Usage
Defining a model in configuration file
A model can be created when invoking a command via the CLI by setting architecture in the configuration file. The task will use the supplied name to get the registered model and instantiate with the provided parameters.
Given the following configuration file configuration.json:
{
...
"architecture:" {
"name": "tcn",
"params": {
"input_kernel": [1, 3],
"input_norm": "batch",
"blocks": [
{"depth": 1, "branch": 1, "filters": 12, "kernel": [1, 3], "dilation": [1, 1], "dropout": 0.10, "ex_ratio": 1, "se_ratio": 0, "norm": "batch"},
{"depth": 1, "branch": 1, "filters": 20, "kernel": [1, 3], "dilation": [1, 1], "dropout": 0.10, "ex_ratio": 1, "se_ratio": 2, "norm": "batch"},
{"depth": 1, "branch": 1, "filters": 28, "kernel": [1, 3], "dilation": [1, 2], "dropout": 0.10, "ex_ratio": 1, "se_ratio": 2, "norm": "batch"},
{"depth": 1, "branch": 1, "filters": 36, "kernel": [1, 3], "dilation": [1, 4], "dropout": 0.10, "ex_ratio": 1, "se_ratio": 2, "norm": "batch"},
{"depth": 1, "branch": 1, "filters": 40, "kernel": [1, 3], "dilation": [1, 8], "dropout": 0.10, "ex_ratio": 1, "se_ratio": 2, "norm": "batch"}
],
"output_kernel": [1, 3],
"include_top": true,
"use_logits": true,
"model_name": "tcn"
}
}
}
The model can be then be created using the following command:
Defining a model in code
Alternatively, the model can be created directly in code using the following snippet: