Skip to content

Model Exporting

Introduction

Export mode is used to convert the trained TensorFlow model into a format that can be used for deployment onto Ambiq's family of SoCs. Currently, the command will convert the TensorFlow model into both TensorFlow Lite (TFL) and TensorFlow Lite for micro-controller (TFLM) variants. The command will also verify the models' outputs match. The activations and weights can be quantized by configuring the quantization section in the configuration file or by setting the quantization parameter in the code.

  1. Load the configuration data (e.g. configuration.json (1))
  2. Load the test data (e.g. test.pkl)
  3. Load the trained model (e.g. model.keras)
  4. Quantize the model (e.g. 16x8)
  5. Convert the model (e.g. TFL, TFLM)
  6. Verify the models' outputs match
  7. Save artifacts (e.g. model.tflite)
  1. Example configuration: 
    {
    "name": "arr-2-eff-sm",
    "project": "hk-rhythm-2",
    "job_dir": "./results/arr-2-eff-sm",
    "verbose": 2,
    "datasets": [
        {
            "name": "ptbxl",
            "params": {
                "path": "./datasets/ptbxl"
            }
        }
    ],
    "num_classes": 2,
    "class_map": {
        "0": 0,
        "7": 1,
        "8": 1
    },
    "class_names": [
        "NORMAL",
        "AFIB/AFL"
    ],
    "class_weights": "balanced",
    "sampling_rate": 100,
    "frame_size": 512,
    "samples_per_patient": [
        10,
        10
    ],
    "val_samples_per_patient": [
        5,
        5
    ],
    "test_samples_per_patient": [
        5,
        5
    ],
    "val_patients": 0.2,
    "val_size": 20000,
    "test_size": 20000,
    "batch_size": 256,
    "buffer_size": 20000,
    "epochs": 100,
    "steps_per_epoch": 50,
    "val_metric": "loss",
    "lr_rate": 0.001,
    "lr_cycles": 1,
    "threshold": 0.75,
    "val_metric_threshold": 0.98,
    "tflm_var_name": "g_rhythm_model",
    "tflm_file": "rhythm_model_buffer.h",
    "backend": "pc",
    "demo_size": 896,
    "display_report": true,
    "quantization": {
        "qat": false,
        "format": "INT8",
        "io_type": "int8",
        "conversion": "CONCRETE",
        "debug": false
    },
    "preprocesses": [
        {
            "name": "layer_norm",
            "params": {
                "epsilon": 0.01,
                "name": "znorm"
            }
        }
    ],
    "augmentations": [],
    "model_file": "model.keras",
    "use_logits": false,
    "architecture": {
        "name": "efficientnetv2",
        "params": {
            "input_filters": 16,
            "input_kernel_size": [
                1,
                9
            ],
            "input_strides": [
                1,
                2
            ],
            "blocks": [
                {
                    "filters": 24,
                    "depth": 2,
                    "kernel_size": [
                        1,
                        9
                    ],
                    "strides": [
                        1,
                        2
                    ],
                    "ex_ratio": 1,
                    "se_ratio": 2
                },
                {
                    "filters": 32,
                    "depth": 2,
                    "kernel_size": [
                        1,
                        9
                    ],
                    "strides": [
                        1,
                        2
                    ],
                    "ex_ratio": 1,
                    "se_ratio": 2
                },
                {
                    "filters": 40,
                    "depth": 2,
                    "kernel_size": [
                        1,
                        9
                    ],
                    "strides": [
                        1,
                        2
                    ],
                    "ex_ratio": 1,
                    "se_ratio": 2
                },
                {
                    "filters": 48,
                    "depth": 1,
                    "kernel_size": [
                        1,
                        9
                    ],
                    "strides": [
                        1,
                        2
                    ],
                    "ex_ratio": 1,
                    "se_ratio": 2
                }
            ],
            "output_filters": 0,
            "include_top": true,
            "use_logits": true
        }
    }
}

graph LR
A("`Load
configuration
__HKTaskParams__
`")
B("`Load test
data
`")
C("`Load trained
__model__
`")
D("`Quantize
__model__
`")
E("`Convert
__model__
`")
F("`Verify
__outputs__
`")
G("`Save
__artifacts__
`")
A ==> B
B ==> C
subgraph CF["Export"]
    C ==> D
    D ==> E
    E ==> F
end
F ==> G

Usage

CLI

The following command will export a rhythm model using the reference configuration.

heartkit --task rhythm --mode export --config ./configuration.json

Python

The model can be evaluated using the following snippet:

1
2
3
4
5
task = hk.TaskFactory.get("rhythm")

params = hk.HKTaskParams(...)  # (1)

task.export(params)
  1. Example configuration: 
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
    hk.HKTaskParams(
    name="arr-2-eff-sm",
    project="hk-rhythm-2",
    job_dir="./results/arr-2-eff-sm",
    verbose=2,
    datasets=[hk.NamedParams(
        name="ptbxl",
        params=dict(
            path="./datasets/ptbxl"
        )
    )],
    num_classes=2,
    class_map={
        "0": 0,
        "7": 1,
        "8": 1
    },
    class_names=[
        "NORMAL", "AFIB/AFL"
    ],
    class_weights="balanced",
    sampling_rate=100,
    frame_size=512,
    samples_per_patient=[10, 10],
    val_samples_per_patient=[5, 5],
    test_samples_per_patient=[5, 5],
    val_patients=0.20,
    val_size=20000,
    test_size=20000,
    batch_size=256,
    buffer_size=20000,
    epochs=100,
    steps_per_epoch=50,
    val_metric="loss",
    lr_rate=1e-3,
    lr_cycles=1,
    threshold=0.75,
    val_metric_threshold=0.98,
    tflm_var_name="g_rhythm_model",
    tflm_file="rhythm_model_buffer.h",
    backend="pc",
    demo_size=896,
    display_report=True,
    quantization=hk.QuantizationParams(
        qat=False,
        format="INT8",
        io_type="int8",
        conversion="CONCRETE",
        debug=False
    ),
    preprocesses=[
        hk.NamedParams(
            name="layer_norm",
            params=dict(
                epsilon=0.01,
                name="znorm"
            )
        )
    ],
    augmentations=[
    ],
    model_file="model.keras",
    use_logits=False,
    architecture=hk.NamedParams(
        name="efficientnetv2",
        params=dict(
            input_filters=16,
            input_kernel_size=[1, 9],
            input_strides=[1, 2],
            blocks=[
                {"filters": 24, "depth": 2, "kernel_size": [1, 9], "strides": [1, 2], "ex_ratio": 1,  "se_ratio": 2},
                {"filters": 32, "depth": 2, "kernel_size": [1, 9], "strides": [1, 2], "ex_ratio": 1,  "se_ratio": 2},
                {"filters": 40, "depth": 2, "kernel_size": [1, 9], "strides": [1, 2], "ex_ratio": 1,  "se_ratio": 2},
                {"filters": 48, "depth": 1, "kernel_size": [1, 9], "strides": [1, 2], "ex_ratio": 1,  "se_ratio": 2}
            ],
            output_filters=0,
            include_top=True,
            use_logits=True
        )
    }
)

Arguments

Please refer to HKTaskParams for the list of arguments that can be used with the export command.