Skip to content

flops

FLOPs Metrics API

FLOPS calculation for keras.Model or keras.Sequential.

Functions:

  • get_flops –

    Calculate FLOPS for keras.Model or keras.Sequential

Functions

get_flops 

get_flops(model: keras.Model, batch_size: int | None = None, fpath: os.PathLike | None = None) -> float

Calculate FLOPS for keras.Model or keras.Sequential. Ignore operations used in only training mode such as Initialization.

Note

Only tensorflow backend is supported currently. Does not support LSTM and GRU.

Parameters:

  • model (Model | Sequential) –

    Model

  • batch_size (int, default: None ) –

    Batch size. Defaults to None.

  • fpath (PathLike, default: None ) –

    Output file path. Defaults to None.

Returns:

  • float ( float ) –

    FLOPS

Example:

model = keras.Sequential([
    keras.layers.Dense(128, activation="relu", input_shape=(784,)),
    keras.layers.Dense(64, activation="relu"),
    keras.layers.Dense(10, activation="softmax"),
])
model.build()
flops = nse.metrics.get_flops(model, batch_size=1)
print(f"FLOPS: {flops/1e6:.2f}M")
Source code in neuralspot_edge/metrics/flops.py 
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
def get_flops(model: keras.Model, batch_size: int | None = None, fpath: os.PathLike | None = None) -> float:
    """Calculate FLOPS for keras.Model or keras.Sequential.
    Ignore operations used in only training mode such as Initialization.

    Note:
        Only tensorflow backend is supported currently.
        Does not support LSTM and GRU.

    Args:
        model (keras.Model|keras.Sequential): Model
        batch_size (int, optional): Batch size. Defaults to None.
        fpath (os.PathLike, optional): Output file path. Defaults to None.

    Returns:
        float: FLOPS

    Example:

    ```python
    model = keras.Sequential([
        keras.layers.Dense(128, activation="relu", input_shape=(784,)),
        keras.layers.Dense(64, activation="relu"),
        keras.layers.Dense(10, activation="softmax"),
    ])
    model.build()
    flops = nse.metrics.get_flops(model, batch_size=1)
    print(f"FLOPS: {flops/1e6:.2f}M")
    ```
    """
    if keras.backend.backend() != "tensorflow":
        raise ValueError("Only tensorflow backend is supported.")

    import tensorflow as tf

    # pylint: disable=no-name-in-module
    from tensorflow.python.profiler.model_analyzer import profile

    # pylint: disable=no-name-in-module
    from tensorflow.python.profiler.option_builder import ProfileOptionBuilder

    input_signature = [tf.TensorSpec(shape=(batch_size,) + model.input_shape[1:])]
    forward_pass = tf.function(model.call, input_signature=input_signature)
    graph = forward_pass.get_concrete_function().graph
    options = ProfileOptionBuilder.float_operation()
    if fpath:
        options["output"] = f"file:outfile={fpath}"
    graph_info = profile(graph, options=options)
    return float(graph_info.total_float_ops)