| 5 | |
| 6 | |
| 7 | def benchmark_concat(num_inputs, input_dim, axis, add_axis, iterations): |
| 8 | input_names = [f"input{i}" for i in range(num_inputs)] |
| 9 | for n in input_names: |
| 10 | workspace.FeedBlob(n, np.random.randn(*input_dim).astype(np.float32)) |
| 11 | |
| 12 | net = core.Net("benchmark_net") |
| 13 | net.Concat(input_names, ["output", "split_info"], axis=axis, add_axis=add_axis) |
| 14 | workspace.CreateNet(net) |
| 15 | |
| 16 | runtimes = workspace.BenchmarkNet(net.Name(), 1, iterations, True) |
| 17 | print(f"{num_inputs * np.prod(input_dim) * 4 / runtimes[1] / 1e6} GB/s") |
| 18 | |
| 19 | |
| 20 | if __name__ == "__main__": |