benchmark reshaping layers to run benchmarks see the following command for an example please change the flag to your custom value python3 m benchmarks layerbenchmark reshapingbenchmark benchmarknamebenchmarkcropping2d numsamples2048 batchsize256 jitcompiletrue this code is a benchmarking tool for various preprocessing and utility layers in neural networks using keras and tensorflow it includes benchmark functions for layers such as cropping1d cropping2d cropping3d flatten permute up sampling1d up sampling2d up sampling3d zero padding1d zero padding2d and zero padding3d each function initializes a layerbenchmark object with the respective layer name and necessary initialization arguments along with the input shape and jit compile option the benchmark predict and benchmark train methods are then used to evaluate the performance of each layer in terms of prediction and training the main function uses the absl library to parse command line arguments allowing the user to specify the benchmark name number of samples batch size and whether to use jit compilation it executes the specified benchmark or all benchmarks if none is specified and handles cases where an invalid benchmark name is provided the benchmarks are stored in a dictionary for easy access and execution
from absl import app
from absl import flags

from benchmarks.layer_benchmark.base_benchmark import LayerBenchmark

FLAGS = flags.FLAGS


def benchmark_cropping1d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Cropping1D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[1024, 256],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_cropping2d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Cropping2D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 256, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_cropping3d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Cropping3D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[32, 32, 32, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_flatten(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Flatten"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 256, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_permute(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Permute"
    init_args = {
        "dims": (2, 1),
    }
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 256],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_up_sampling1d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "UpSampling1D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_up_sampling2d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "UpSampling2D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[128, 128, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_up_sampling3d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "UpSampling3D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[32, 16, 16, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_zero_padding1d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "ZeroPadding1D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_zero_padding2d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "ZeroPadding2D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[256, 256, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


def benchmark_zero_padding3d(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "ZeroPadding3D"
    init_args = {}
    benchmark = LayerBenchmark(
        layer_name,
        init_args,
        input_shape=[32, 32, 32, 3],
        jit_compile=jit_compile,
    )

    benchmark.benchmark_predict(
        num_samples=num_samples,
        batch_size=batch_size,
    )

    benchmark.benchmark_train(
        num_samples=num_samples,
        batch_size=batch_size,
    )


BENCHMARK_NAMES = {
    "benchmark_cropping1d": benchmark_cropping1d,
    "benchmark_cropping2d": benchmark_cropping2d,
    "benchmark_cropping3d": benchmark_cropping3d,
    "benchmark_flatten": benchmark_flatten,
    "benchmark_permute": benchmark_permute,
    "benchmark_up_sampling1d": benchmark_up_sampling1d,
    "benchmark_up_sampling2d": benchmark_up_sampling2d,
    "benchmark_up_sampling3d": benchmark_up_sampling3d,
    "benchmark_zero_padding1d": benchmark_zero_padding1d,
    "benchmark_zero_padding2d": benchmark_zero_padding2d,
    "benchmark_zero_padding3d": benchmark_zero_padding3d,
}


def main(_):
    benchmark_name = FLAGS.benchmark_name
    num_samples = FLAGS.num_samples
    batch_size = FLAGS.batch_size
    jit_compile = FLAGS.jit_compile

    if benchmark_name is None:
        for name, benchmark_fn in BENCHMARK_NAMES.items():
            benchmark_fn(num_samples, batch_size, jit_compile)
        return

    if benchmark_name not in BENCHMARK_NAMES:
        raise ValueError(
            f"Invalid benchmark name: {benchmark_name}, `benchmark_name` must "
            f"be one of {BENCHMARK_NAMES.keys()}"
        )
    benchmark_fn = BENCHMARK_NAMES[benchmark_name]
    benchmark_fn(num_samples, batch_size, jit_compile)


if __name__ == "__main__":
    app.run(main)