benchmark activation layers to run benchmarks see the following command for an example please change the flag to your custom value python3 m benchmarks layerbenchmark activationbenchmark benchmarknamebenchmarkelu numsamples2048 batchsize256 jitcompile true this code defines a set of functions to benchmark different neural network layers including elu prelu relu leaky relu and softmax each function takes the number of samples batch size and a boolean for jit compile as arguments and creates a layer benchmark object with specified layer name empty initialization arguments a fixed input shape and the jit compile option it then runs prediction and training benchmarks on this layer the main function parses command line arguments for the benchmark name number of samples batch size and jit compile option it runs the specified benchmark or all benchmarks if none is specified using the provided command line arguments the code also includes a dictionary mapping benchmark names to their respective functions and uses the absl library for handling command line arguments and running the main function
from absl import app
from absl import flags

from benchmarks.layer_benchmark.base_benchmark import LayerBenchmark

FLAGS = flags.FLAGS


def benchmark_elu(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "ELU"
    init_args = {}
    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_prelu(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "PReLU"
    init_args = {}
    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_relu(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "ReLU"
    init_args = {}
    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_leaky_relu(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "LeakyReLU"
    init_args = {}
    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_softmax(
    num_samples,
    batch_size,
    jit_compile=True,
):
    layer_name = "Softmax"
    init_args = {}
    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,
    )


BENCHMARK_NAMES = {
    "benchmark_elu": benchmark_elu,
    "benchmark_relu": benchmark_relu,
    "benchmark_leaky_relu": benchmark_leaky_relu,
    "benchmark_prelu": benchmark_prelu,
    "benchmark_softmax": benchmark_softmax,
}


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)