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
|
# SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Rewrite functions used to return layers ready for sparse pruning."""
from typing import Any
import tensorflow_model_optimization as tfmot
from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.rewrite.library.helper_functions import compute_conv2d_parameters
from mlia.nn.rewrite.library.helper_functions import get_activation_function
def fc_sparsity_rewrite(
input_shape: Any, output_shape: Any, sparsity_m: int = 2, sparsity_n: int = 4
) -> keras.Model:
"""Fully connected TensorFlow Lite model ready for sparse pruning."""
model = tfmot.sparsity.keras.prune_low_magnitude(
to_prune=keras.Sequential(
[
keras.layers.InputLayer(input_shape=input_shape),
keras.layers.Reshape([-1]),
keras.layers.Dense(output_shape),
]
),
sparsity_m_by_n=(
sparsity_m,
sparsity_n,
),
)
return model
def conv2d_sparsity_rewrite( # pylint: disable=dangerous-default-value
input_shape: Any,
output_shape: Any,
sparsity_m: int = 2,
sparsity_n: int = 4,
activation: str = "relu",
kernel_size: list[int] = [3, 3],
) -> keras.Model:
"""Conv2d TensorFlow Lite model ready for sparse pruning."""
conv2d_parameters = compute_conv2d_parameters(
input_shape=input_shape,
output_shape=output_shape,
kernel_size_input=kernel_size,
)
activation_function, activation_function_extra_args = get_activation_function(
activation
)
activation_func_found = (
[activation_function(**activation_function_extra_args)]
if activation_function
else []
)
model = tfmot.sparsity.keras.prune_low_magnitude(
to_prune=keras.Sequential(
[
keras.layers.InputLayer(input_shape=input_shape),
keras.layers.Conv2D(**conv2d_parameters),
keras.layers.BatchNormalization(),
*activation_func_found,
]
),
sparsity_m_by_n=(
sparsity_m,
sparsity_n,
),
)
return model
|
