| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import time |
| import functools |
| import numpy as np |
| |
| from .common_utils import merge_map_list |
| from mxnet import profiler |
| |
| """ |
| TODO: Below we are using logic of parsing the MXNet profiler output string to |
| fetch the benchmark results. Note that this is a temporary solution till we add |
| a new utility API into MXNet profiler to get_summary(), reset(). All the below |
| parsing logic should be removed once these read APIs are available in Profiler. |
| |
| """ |
| |
| |
| def _get_memory_profile(memory_profile_results): |
| memory_profile = {} |
| for line in memory_profile_results: |
| if line.startswith("Memory:"): |
| device_id = line.split()[1] |
| avg_time_memory_alloc = float(line.split()[-1]) |
| memory_profile["max_storage_mem_alloc_" + device_id] = avg_time_memory_alloc |
| |
| return memory_profile |
| |
| |
| def _get_operator_profile(operator_name, operator_profile_results): |
| operator_profile = {} |
| |
| # alias map : dictionary of the form {"alias" : "registered_name"} |
| # allows to retrieve alias operator profile from the profiler results |
| alias_map = {"broadcast_plus": "broadcast_add", "broadcast_minus": "broadcast_sub", "flatten": "Flatten", "max_axis": "max", "Custom": "CustomAddOne", |
| "swapaxes": "SwapAxis", "flip": "reverse", "reshape": "Reshape", "crop": "slice", "sum_axis": "sum", "min_axis": "min", "ctc_loss": "CTCLoss", |
| "fill_element_0index": "TernaryOp", "identity": "_copy", "ElementWiseSum": "add_n", "choose_element_0index": "pick", "stop_gradient": "BlockGrad", |
| "broadcast_axes": "broadcast_axis"} |
| |
| op_name = None |
| |
| if operator_name in alias_map: |
| op_name = alias_map[operator_name] |
| else: |
| op_name = operator_name |
| |
| # Variables to store forward/backward performance results |
| forward_res, backward_res = None, None |
| |
| for line in operator_profile_results: |
| if op_name in line or op_name[:3] + " " in line: |
| operation = line.split()[0] |
| operation_avg_time = float(line.split()[-1]) |
| if "_backward" in operation: |
| backward_res = operation_avg_time |
| else: |
| forward_res = operation_avg_time |
| |
| # Add forward and backward performance results to the dict in the correct order |
| if forward_res: |
| operator_profile["avg_time_forward_" + operator_name] = forward_res |
| |
| if backward_res: |
| operator_profile["avg_time_backward_" + operator_name] = backward_res |
| |
| return operator_profile |
| |
| |
| def parse_profiler_dump(operator_name, profiler_dump): |
| """Parse the MXNet profiler dump output, fetch Memory profile results and |
| Operator compute profiler results. |
| |
| Parameters |
| ---------- |
| profiler_dump: string |
| MXNet profiler output from mx.profiler.dumps() API. |
| |
| Returns |
| ------- |
| map, Memory and Compute profiler results. |
| |
| """ |
| if not profiler_dump: |
| raise AssertionError("Invalid MXNet profiler output provided to parse!") |
| |
| """ |
| MXNet profiler output from mx.profiler.dumps() API looks like below. This function parses |
| this string profiler output to fetch Memory and Compute metrics. |
| |
| Profile Statistics. |
| Note that counter items are counter values and not time units. |
| Device Storage |
| ================= |
| Name Total Count Time (ms) Min Time (ms) Max Time (ms) Avg Time (ms) |
| ---- ----------- --------- ------------- ------------- ------------- |
| Memory: cpu/0 100 2097152.0000 1681915.8750 2097152.0000 207618.0469 |
| |
| MXNET_C_API |
| ================= |
| Name Total Count Time (ms) Min Time (ms) Max Time (ms) Avg Time (ms) |
| ---- ----------- --------- ------------- ------------- ------------- |
| MXNDArrayFree 49 1.1220 0.0170 0.0360 0.0229 |
| MXAutogradBackwardEx 50 11.5460 0.1980 0.3360 0.2309 |
| MXNet C API Calls 399 1.9990 1.6010 1.9990 0.1990 |
| MXImperativeInvoke 50 4.4810 0.0700 0.1330 0.0896 |
| MXNDArrayWaitAll 50 769.0570 14.0200 24.5030 15.3811 |
| MXAutogradSetIsTraining 100 0.0190 0.0000 0.0010 0.0002 |
| MXAutogradSetIsRecording 100 0.0400 0.0000 0.0010 0.0004 |
| MXNet C API Concurrency 798 0.0000 0.0000 0.0010 0.0005 |
| |
| operator |
| ================= |
| Name Total Count Time (ms) Min Time (ms) Max Time (ms) Avg Time (ms) |
| ---- ----------- --------- ------------- ------------- ------------- |
| DeleteVariable 196 1.4490 0.0040 0.0250 0.0074 |
| _backward_broadcast_add 100 521.2320 4.8070 8.5970 5.2123 |
| SetValueOp 100 645.8060 5.8820 10.0380 6.4581 |
| broadcast_add 100 394.8910 3.5230 5.8790 3.9489 |
| """ |
| |
| # String Patterns to look out for when parsing |
| memory_profile_result_start = "Device Storage" # Helps identify start of Memory profile |
| c_api_profile_result_start = "MXNET_C_API" # Helps identify end of Memory profile |
| |
| if operator_name == "Custom": |
| operator_profile_result_start = "Custom Operator" # Helps identify start of Custom Operator profile |
| else: |
| operator_profile_result_start = "operator" # Helps identify start of Operator profile |
| |
| memory_profile_results = [] |
| operator_profile_results = [] |
| |
| # Parse lines corresponding to Memory and Computation profiling |
| read_memory_profile = False |
| read_operator_profile = False |
| for line in profiler_dump.splitlines(): |
| if line.startswith(memory_profile_result_start): |
| read_memory_profile = True |
| elif line.startswith(operator_profile_result_start): |
| read_operator_profile = True |
| elif line.startswith(c_api_profile_result_start): |
| read_memory_profile = False |
| |
| if read_memory_profile: |
| memory_profile_results.append(line) |
| elif read_operator_profile: |
| operator_profile_results.append(line) |
| |
| # Prepare results |
| memory_profile = _get_memory_profile(memory_profile_results) |
| operator_profile = _get_operator_profile(operator_name, operator_profile_results) |
| |
| return merge_map_list([memory_profile, operator_profile]) |
| |
| |
| def cpp_profile(func): |
| """Decorator for profiling MXNet operation. |
| Uses MXNet profiler to collect metrics on memory usage and execution time |
| of the operation. |
| |
| Parameters |
| ---------- |
| func: |
| Operation to be executed and timed. |
| |
| Returns |
| ------- |
| res, profiler output. res being result returned after operator execution. |
| profiler output is a dictionary with summary of operation execution. |
| Example output : { "add": [{"avg_time_mem_alloc_cpu/0": 207618.0469, |
| "avg_time_forward_broadcast_add": 4.204, |
| "avg_time_backward_broadcast_add": 5.6288, |
| "inputs": { |
| "lhs": [1024, 1024], |
| "rhs": [1024,1024] |
| }] |
| } |
| """ |
| |
| @functools.wraps(func) |
| def cpp_profile_it(*args, **kwargs): |
| # Profile the operation |
| profiler.set_config(profile_all=True, aggregate_stats=True) |
| profiler.set_state('run') |
| res = func(*args, **kwargs) |
| profiler.set_state('stop') |
| |
| # Prepare the results |
| profiler_dump = profiler.dumps(reset=True) |
| |
| # args[0] is assumed to be operator name, if not found check for block name. |
| # NOTE: This parameter should be removed when we get away from parsing |
| # profiler output and start using new profiler APIs - get_summary(), reset() |
| if len(args) > 0: |
| operator_name = args[0].__name__ |
| elif 'block' in kwargs: |
| operator_name = kwargs['block']._op_name |
| else: |
| raise ValueError("Unable to identify operator name to extract profiler output!") |
| |
| # Get the MXNet profile output |
| profiler_output = parse_profiler_dump(operator_name, profiler_dump) |
| return res, profiler_output |
| |
| return cpp_profile_it |
| |
| |
| def python_profile(func): |
| """Decorator for profiling MXNet operation. |
| Uses Python's time module to collect execution time information |
| of the operation. |
| |
| Parameters |
| ---------- |
| func: |
| Operation to be executed and timed. |
| |
| Returns |
| ------- |
| res, timing output. res being result returned after operator execution. |
| profiler output is a dictionary with summary of operation execution. |
| Example output : { "add": [{"avg_time_add": 0.4053089120425284, |
| 'p50_time_add': 16.761042876169086, |
| 'p90_time_add': 18.081666342914108, |
| 'p99_time_add': 19.060144051909447, |
| "inputs": { |
| "lhs": [1024, 1024], |
| "rhs": [1024,1024] |
| }] |
| } |
| """ |
| |
| @functools.wraps(func) |
| def python_profile_it(*args, **kwargs): |
| runs = args[1] |
| modified_args = (args[0], 1) |
| times = [] |
| |
| for _ in range(runs): |
| start_time = time.perf_counter() # 1 |
| res = func(*modified_args, **kwargs) |
| end_time = time.perf_counter() # 2 |
| run_time = (end_time - start_time)*1000 # 3 |
| times.append(run_time) |
| |
| # NOTE : same as cpp_profile_it |
| if len(args) > 0: |
| operator_name = args[0].__name__ |
| elif 'block' in kwargs: |
| operator_name = kwargs['block']._op_name |
| else: |
| raise ValueError("Unable to identify operator name to extract profiler output!") |
| |
| avg_run_time = np.mean(times) |
| p50_run_time = np.percentile(times, 50) |
| p90_run_time = np.percentile(times, 90) |
| p99_run_time = np.percentile(times, 99) |
| |
| profiler_output = {'avg_time_'+str(operator_name): avg_run_time, |
| 'p50_time_'+str(operator_name): p50_run_time, |
| 'p90_time_'+str(operator_name): p90_run_time, |
| 'p99_time_'+str(operator_name): p99_run_time, |
| } |
| return res, profiler_output |
| return python_profile_it |
