337 lines
7.0 KiB
Python
337 lines
7.0 KiB
Python
""" Benchmark reshaping layers.
|
|
|
|
To run benchmarks, see the following command for an example, please change the
|
|
flag to your custom value:
|
|
|
|
```
|
|
python3 -m benchmarks.layer_benchmark.reshaping_benchmark \
|
|
--benchmark_name=benchmark_cropping2d \
|
|
--num_samples=2048 \
|
|
--batch_size=256 \
|
|
--jit_compile=True
|
|
```
|
|
"""
|
|
|
|
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=[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_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 = {
|
|
"dim": (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=[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_up_sampling3d(
|
|
num_samples,
|
|
batch_size,
|
|
jit_compile=True,
|
|
):
|
|
layer_name = "UpSampling3D"
|
|
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_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:
|
|
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)
|