1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18import mxnet as mx
19from mxnet import nd
20
21from benchmark.opperf.utils.benchmark_utils import run_performance_test
22from benchmark.opperf.utils.benchmark_utils import run_op_benchmarks
23from benchmark.opperf.utils.op_registry_utils import get_all_optimizer_operators
24from benchmark.opperf.utils.common_utils import merge_map_list
25from benchmark.opperf.rules.default_params import MX_OP_MODULE
26
27"""Performance benchmark tests for MXNet Neural Network Optimizer Update Operators.
28
291. Stochastic Gradient Descent (SGD)
30    1.1 mp_sgd_update
31    1.2 sgd_mom_update
32    1.3 signsgd_update
33    1.4 mp_sgd_mom_update
34    1.5 sgd_update
352. signum_update
363. rmspropalex_update
374. ftml_update
385. rmsprop_update
396. ftrl_update
407. adam_update
418. preloaded_multi_*
42    8.1 preloaded_multi_sgd_mom_update
43    8.2 preloaded_multi_sgd_update
44    8.3 preloaded_multi_mp_sgd_update
45    8.4 preloaded_multi_mp_sgd_mom_update
469. lamb_*
47    9.1 lamb_update_phase1
48    9.2 lamb_update_phase2
4910. multi_*
50    10.1 multi_sgd_update
51    10.2 multi_sgd_mom_update
52    10.3 multi_mp_sgd_update
53    10.4 multi_mp_sgd_mom_update
54"""
55
56
57def run_optimizer_operators_benchmarks(ctx=mx.cpu(), dtype='float32', profiler='native', warmup=25, runs=100):
58    """Runs benchmarks with the given context and precision (dtype) for all the neural network
59    optimizer update operators in MXNet.
60
61    Parameters
62    ----------
63    ctx: mx.ctx
64        Context to run benchmarks
65    dtype: str, default 'float32'
66        Precision to use for benchmarks
67    profiler: str, default 'native'
68        Type of Profiler to use (native/python)
69    warmup: int, default 25
70        Number of times to run for warmup
71    runs: int, default 100
72        Number of runs to capture benchmark results
73
74    Returns
75    -------
76    Dictionary of results. Key -> Name of the operator, Value -> Benchmark results.
77
78    """
79    # Run independent tests for ops that need specific input data
80    multi_mp_sgd_mom_res = run_performance_test([getattr(MX_OP_MODULE, "multi_mp_sgd_mom_update")],
81                                                inputs=[{"args0": nd.random_normal(shape=(5,5)),
82                                                "args1": nd.random_normal(shape=(5,5)), "args2": nd.random_normal(shape=(5,5)),
83                                                "args3": nd.random_normal(shape=(5,5)), "lrs": 0.1, "wds": 0.2,
84                                                "out": nd.random_normal(shape=(5,5))}],run_backward=False)
85
86    multi_sgd_mom_res = run_performance_test([getattr(MX_OP_MODULE, "multi_sgd_mom_update")],
87                                             inputs=[{"args0": nd.random_normal(shape=(5,5)),
88                                             "args1": nd.random_normal(shape=(5,5)),"args2": nd.random_normal(shape=(5,5)),
89                                             "lrs": 0.1, "wds": 0.2, "out": nd.random_normal(shape=(5,5))}], run_backward=False)
90
91    multi_sgd_res = run_performance_test([getattr(MX_OP_MODULE, "multi_sgd_update")],
92                                         inputs=[{"args0": nd.random_normal(shape=(5,5)),
93                                         "args1": nd.random_normal(shape=(5,5)), "lrs": 0.1, "wds": 0.2,
94                                         "out": nd.random_normal(shape=(5,5))}], run_backward=False)
95
96    multi_mp_sgd_res = run_performance_test([getattr(MX_OP_MODULE, "multi_mp_sgd_update")],
97                                            inputs=[{"args0": nd.random_normal(shape=(5,5)),
98                                            "args1": nd.random_normal(shape=(5,5)),"args2": nd.random_normal(shape=(5,5)),
99                                            "lrs": 0.1, "wds": 0.2, "out": nd.random_normal(shape=(5,5))}], run_backward=False)
100
101    preloaded_multi_mp_sgd_res = run_performance_test(
102                                 [getattr(MX_OP_MODULE, "preloaded_multi_mp_sgd_update")],
103                                 inputs=[{"args0": nd.random_normal(shape=(5,5)),
104                                          "args1": nd.random_normal(shape=(5,5)), "args2": nd.random_normal(shape=(5,5)),
105                                          "args3": nd.random_normal(shape=(1)), "args4": nd.random_normal(shape=(1)),
106                                          "out": nd.random_normal(shape=(5,5))}], run_backward=False)
107
108    preloaded_multi_sgd_mom_res = run_performance_test(
109                                  [getattr(MX_OP_MODULE, "preloaded_multi_sgd_mom_update")],
110                                  inputs=[{"args0": nd.random_normal(shape=(5,5)),
111                                           "args1": nd.random_normal(shape=(5,5)), "args2": nd.random_normal(shape=(5,5)),
112                                           "args3": nd.random_normal(shape=(1)), "args4": nd.random_normal(shape=(1)),
113                                           "out": nd.random_normal(shape=(5,5))}], run_backward=False)
114
115    preloaded_multi_sgd_res = run_performance_test(
116                              [getattr(MX_OP_MODULE, "preloaded_multi_sgd_update")],
117                              inputs=[{"args0": nd.random_normal(shape=(5,5)), "args1": nd.random_normal(shape=(5,5)),
118                                       "args4": nd.random_normal(shape=(1)), "args5": nd.random_normal(shape=(1)),
119                                       "out": nd.random_normal(shape=(5,5))}], run_backward=False)
120
121    preloaded_multi_mp_sgd_mom_res = run_performance_test(
122                                     [getattr(MX_OP_MODULE, "preloaded_multi_mp_sgd_mom_update")],
123                                     inputs=[{"args0": nd.random_normal(shape=(5,5)), "args1": nd.random_normal(shape=(5,5)),
124                                              "args2": nd.random_normal(shape=(5,5)), "args3": nd.random_normal(shape=(5,5)),
125                                              "args4": nd.random_normal(shape=(1)), "args5": nd.random_normal(shape=(1)),
126                                              "out": nd.random_normal(shape=(5,5))}], run_backward=False)
127
128    # Fetch remaining optimizer operators
129    mx_optimizer_ops = get_all_optimizer_operators()
130
131    # Run benchmarks
132    mx_optimizer_op_results = run_op_benchmarks(mx_optimizer_ops, dtype, ctx, profiler, warmup, runs)
133    return merge_map_list(multi_sgd_mom_res + multi_sgd_mom_res + multi_sgd_res + multi_mp_sgd_res + preloaded_multi_mp_sgd_res +\
134                          preloaded_multi_sgd_mom_res + preloaded_multi_mp_sgd_res + preloaded_multi_mp_sgd_mom_res +\
135                          [mx_optimizer_op_results])
136