Performance issue when using sionna.phy.Block.build in graph mode

[jasperbrg]

Hello,

I ran into a significant performance issue in case of involved computations in the build routine of a custom Block.

I created the following minimal example below, where I compare the performance of a Sionna Block and a Keras Layer. Both execute the same function build_routine() in their build. Meanwhile, build_routine measures its runtime. I am calling build_routine in different ways: standalone, through calling the Block in eager mode, calling the Layer in eager mode and calling Block and Layer in graph mode, respectively.

All calls print roughly the same runtime, except calling Block in graph mode, which results in a runtime approximately 10x higher.

I am using Sionna 1.0.2 with Tensorflow 2.14.0.

Am I making a mistake here? I would expect the Block to perform very similarly to the Layer .

Has anyone had similar issues?
Thanks in advance.

import os
import timeit

import numpy as np
from keras.src.engine.base_layer import Layer
from sionna.phy import Block

import tensorflow as tf

gpus = tf.config.list_physical_devices("GPU")
print("\nNum GPUs Available: ", len(gpus))
gpu_num = 0  # Use "" to use the CPU
if len(gpus) == 0:
    gpu_num = ""
os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_num}"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

if gpus:
    try:
        print("Enabling memory growth.")
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print("Enabling memory growth failed.")
        print(e)


def build_routine():
    start = timeit.default_timer()
    x = np.ones(1000)
    for i in range(2**20):
        x = x.copy()
    stop = timeit.default_timer()
    print("Time: ", stop - start)


class MyBlock(Block):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def build(self, *arg_shapes, **kwarg_shapes):
        build_routine()

    def call(self, *args, **kwargs):
        pass


class MyLayer(Layer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def build(self, input_shape):
        build_routine()

    def call(self, inputs, *args, **kwargs):
        pass


if __name__ == "__main__":

    print("Standalone")
    build_routine()

    block = MyBlock()
    layer = MyLayer()

    print("Block eager")
    block()

    print("Layer eager")
    layer(None)

    block2 = MyBlock()
    layer2 = MyLayer()

    tf_func_block = tf.function(block2, jit_compile=False)
    tf_func_layer = tf.function(layer2, jit_compile=False)

    print("Block graph")
    tf_func_block()

    print("Layer graph")
    tf_func_layer(None)

[jasperbrg]

I narrowed it down a bit more. So the issue seems to stem from the combination of using tf.init_scope within tf.function.

import os
import timeit

import numpy as np

import tensorflow as tf

gpus = tf.config.list_physical_devices("GPU")
print("\nNum GPUs Available: ", len(gpus))
gpu_num = 0  # Use "" to use the CPU
if len(gpus) == 0:
    gpu_num = ""
os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_num}"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

if gpus:
    try:
        print("Enabling memory growth.")
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print("Enabling memory growth failed.")
        print(e)


def build_routine():
    start = timeit.default_timer()
    x = np.ones(1000)
    for i in range(2**20):
        x = x.copy()
    stop = timeit.default_timer()
    print("Time: ", stop - start)


if __name__ == "__main__":
    print("Standalone")
    build_routine()

    print("Init scope")
    with tf.init_scope():
        build_routine()

    @tf.function
    def func():
        with tf.init_scope():
            build_routine()

    print("Graph mode init scope")
    func()

[jhoydis]

Hi @jasperbrg,

Thanks for reporting this issue. I am not sure to have an easy fix for that, but we look into it.

Apart from the fact that the build might run longer, there should be no performance difference, once it is compiled, right?

[jasperbrg]

@jhoydis, thanks for the quick response.

That is right, I did not notice any other difference in performance so far. I am just doing a lot of offline numpy computations during build to design the block, so it is quite significant in my use case.

[jhoydis]

Hm, do you need to do it in the build? If you could infer the shapes already earlier, you could initialize the block in the init function.

[jasperbrg]

Yes, that is what I will probably do as a workaround.