codekingpro/portable-devtools

1import argparse
2import os
3
4import numpy
5import psutil
6from onnx import TensorProto
7
8"""
9This profiler tool could run a transformer model and print out the kernel time spent on each Node of the model.
10Example of profiling of longformer model:
11    python profiler.py --model longformer-base-4096_fp32.onnx --batch_size 1 --sequence_length 4096 --global_length 8 --samples 1000 --thread_num 8 --dummy_inputs longformer --use_gpu
12Example of importing profile result file from onnxruntime_perf_test:
13    python profiler.py --input profile_2021-10-25_12-02-41.json
14"""
15
16
17def parse_arguments(argv=None):
18    parser = argparse.ArgumentParser()
19
20    parser.add_argument(
21        "-i",
22        "--input",
23        required=False,
24        type=str,
25        help="Set the input file for reading the profile results",
26    )
27
28    parser.add_argument(
29        "-m",
30        "--model",
31        required=False,
32        type=str,
33        help="onnx model path to run profiling. Required when --input is not specified.",
34    )
35
36    parser.add_argument(
37        "-b",
38        "--batch_size",
39        required=False,
40        type=int,
41        default=1,
42        help="batch size of input",
43    )
44
45    parser.add_argument(
46        "-s",
47        "--sequence_length",
48        required=False,
49        type=int,
50        default=32,
51        help="sequence length of input",
52    )
53
54    parser.add_argument(
55        "--past_sequence_length",
56        required=False,
57        type=int,
58        default=1,
59        help="past sequence length for gpt2",
60    )
61
62    parser.add_argument(
63        "--global_length",
64        required=False,
65        type=int,
66        default=1,
67        help="number of global tokens for longformer",
68    )
69
70    parser.add_argument(
71        "--samples",
72        required=False,
73        type=int,
74        default=1000,
75        help="number of samples to test. Set it large enough to reduce the variance of performance result.",
76    )
77
78    parser.add_argument(
79        "--threshold",
80        required=False,
81        type=float,
82        default=0.01,
83        help="Threshold of run time ratio among all nodes. Nodes with larger ratio will show in top expensive nodes.",
84    )
85
86    parser.add_argument(
87        "--thread_num",
88        required=False,
89        type=int,
90        default=-1,
91        help="number of threads to use",
92    )
93
94    parser.add_argument(
95        "--input_ids_name",
96        required=False,
97        type=str,
98        default=None,
99        help="input name for input IDs, for bert",
100    )
101    parser.add_argument(
102        "--segment_ids_name",
103        required=False,
104        type=str,
105        default=None,
106        help="input name for segment IDs, for bert",
107    )
108    parser.add_argument(
109        "--input_mask_name",
110        required=False,
111        type=str,
112        default=None,
113        help="input name for attention mask, for bert",
114    )
115
116    parser.add_argument(
117        "--dummy_inputs",
118        required=False,
119        default="default",
120        choices=["bert", "gpt2", "longformer", "default"],
121        help="Type of model inputs. The default will create dummy inputs with ones.",
122    )
123
124    parser.add_argument("-g", "--use_gpu", required=False, action="store_true", help="use GPU")
125    parser.set_defaults(use_gpu=False)
126
127    parser.add_argument(
128        "--provider",
129        required=False,
130        type=str,
131        default="cuda",
132        help="Execution provider to use",
133    )
134
135    parser.add_argument(
136        "--basic_optimization",
137        required=False,
138        action="store_true",
139        help="Enable only basic graph optimizations. By default, all optimizations are enabled in OnnxRuntime",
140    )
141    parser.set_defaults(basic_optimization=False)
142
143    parser.add_argument(
144        "--kernel_time_only",
145        required=False,
146        action="store_true",
147        help="Only include the kernel time and no fence time",
148    )
149    parser.set_defaults(kernel_time_only=False)
150
151    parser.add_argument("-v", "--verbose", required=False, action="store_true")
152    parser.set_defaults(verbose=False)
153
154    return parser.parse_args(argv)
155
156
157def run_profile(onnx_model_path, use_gpu, provider, basic_optimization, thread_num, all_inputs):
158    from benchmark_helper import create_onnxruntime_session  # noqa: PLC0415
159
160    session = create_onnxruntime_session(
161        onnx_model_path,
162        use_gpu,
163        provider,
164        enable_all_optimization=not basic_optimization,
165        num_threads=thread_num,
166        enable_profiling=True,
167    )
168
169    for inputs in all_inputs:
170        _ = session.run(None, inputs)
171
172    profile_file = session.end_profiling()
173    return profile_file
174
175
176def get_dim_from_type_proto(dim):
177    return getattr(dim, dim.WhichOneof("value")) if type(dim.WhichOneof("value")) == str else None  # noqa: E721
178
179
180def get_shape_from_type_proto(type_proto):
181    return [get_dim_from_type_proto(d) for d in type_proto.tensor_type.shape.dim]
182
183
184def create_dummy_inputs(onnx_model, batch_size, sequence_length, samples):
185    """Create dummy inputs for ONNX model.
186
187    Args:
188        onnx_model (OnnxModel): ONNX model
189        batch_size (int): batch size
190        sequence_length (int): sequence length
191        samples (int): number of samples
192
193    Returns:
194        List[Dict]: list of inputs
195    """
196    dummy_inputs = {}
197    for graph_input in onnx_model.get_graph_inputs_excluding_initializers():
198        shape = get_shape_from_type_proto(graph_input.type)
199        symbol_dims = []
200        for i, dim in enumerate(shape):
201            if isinstance(dim, str):
202                symbol_dims.append(i)
203
204        # allowed symbolic dimensions: batch_size and sequence_length
205        if len(symbol_dims) > 2:
206            return None
207        if len(symbol_dims) > 0:
208            shape[symbol_dims[0]] = batch_size
209        if len(symbol_dims) > 1:
210            shape[symbol_dims[1]] = sequence_length
211
212        elem_type = graph_input.type.tensor_type.elem_type
213        assert elem_type in [TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64]
214        data_type = (
215            numpy.float32
216            if elem_type == TensorProto.FLOAT
217            else (numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32)
218        )
219        data = numpy.ones(shape, dtype=data_type)
220        dummy_inputs[graph_input.name] = data
221
222    all_inputs = [dummy_inputs for _ in range(samples)]
223    return all_inputs
224
225
226def create_bert_inputs(
227    onnx_model,
228    batch_size,
229    sequence_length,
230    samples,
231    input_ids_name=None,
232    segment_ids_name=None,
233    input_mask_name=None,
234):
235    """Create dummy inputs for BERT model.
236
237    Args:
238        onnx_model (OnnxModel): ONNX model
239        batch_size (int): batch size
240        sequence_length (int): sequence length
241        samples (int): number of samples
242        input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None.
243        segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None.
244        input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None.
245
246    Returns:
247        List[Dict]: list of inputs
248    """
249    from bert_test_data import find_bert_inputs, generate_test_data  # noqa: PLC0415
250
251    input_ids, segment_ids, input_mask = find_bert_inputs(onnx_model, input_ids_name, segment_ids_name, input_mask_name)
252    all_inputs = generate_test_data(
253        batch_size,
254        sequence_length,
255        test_cases=samples,
256        seed=123,
257        verbose=False,
258        input_ids=input_ids,
259        segment_ids=segment_ids,
260        input_mask=input_mask,
261        random_mask_length=False,
262    )
263
264    return all_inputs
265
266
267def create_gpt2_inputs(onnx_model, batch_size, sequence_length, past_sequence_length, samples):
268    """Create dummy inputs for GPT-2 model.
269
270    Args:
271        onnx_model (OnnxModel): ONNX model
272        batch_size (int): batch size
273        sequence_length (int): sequence length
274        past_sequence_length (int): past sequence length
275        samples (int): number of samples
276
277    Raises:
278        RuntimeError: symbolic is not supported. Use the tool convert_to_onnx.py to export ONNX model instead.
279
280    Returns:
281        List[Dict]: list of inputs
282    """
283    # The symbolic names shall be same as those used in Gpt2Helper.export_onnx(...) function.
284    symbols = {
285        "batch_size": batch_size,
286        "seq_len": sequence_length,
287        "past_seq_len": past_sequence_length,
288        "total_seq_len": sequence_length + past_sequence_length,
289    }
290
291    dummy_inputs = {}
292    for graph_input in onnx_model.get_graph_inputs_excluding_initializers():
293        shape = get_shape_from_type_proto(graph_input.type)
294        for i, dim in enumerate(shape):
295            if isinstance(dim, str):
296                if dim not in symbols:
297                    raise RuntimeError(f"symbol is not supported: {dim}")
298                else:
299                    shape[i] = symbols[dim]
300
301        elem_type = graph_input.type.tensor_type.elem_type
302        assert elem_type in [TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64]
303        data_type = (
304            numpy.float32
305            if elem_type == TensorProto.FLOAT
306            else (numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32)
307        )
308        data = numpy.ones(shape, dtype=data_type)
309        dummy_inputs[graph_input.name] = data
310
311    all_inputs = [dummy_inputs for _ in range(samples)]
312    return all_inputs
313
314
315def create_longformer_inputs(onnx_model, batch_size, sequence_length, global_length, samples):
316    """Create dummy inputs for Longformer model.
317
318    Args:
319        onnx_model (OnnxModel): ONNX model
320        batch_size (int): batch size
321        sequence_length (int): sequence length
322        global_length (int): number of global tokens
323        samples (int): number of samples
324
325    Raises:
326        RuntimeError: symbolic is not supported. Use the tool convert_longformer_to_onnx.py to export ONNX model instead.
327
328    Returns:
329        List[Dict]: list of inputs
330    """
331    symbols = {"batch_size": batch_size, "sequence_length": sequence_length}
332
333    dummy_inputs = {}
334    for graph_input in onnx_model.get_graph_inputs_excluding_initializers():
335        shape = get_shape_from_type_proto(graph_input.type)
336        for i, dim in enumerate(shape):
337            if isinstance(dim, str):
338                if dim not in symbols:
339                    raise RuntimeError(f"symbol is not supported: {dim}")
340                else:
341                    shape[i] = symbols[dim]
342
343        elem_type = graph_input.type.tensor_type.elem_type
344        assert elem_type in [TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64]
345        data_type = (
346            numpy.float32
347            if elem_type == TensorProto.FLOAT
348            else (numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32)
349        )
350
351        if "global" in graph_input.name:
352            data = numpy.zeros(shape, dtype=data_type)
353            data[:, :global_length] = 1
354        else:
355            data = numpy.ones(shape, dtype=data_type)
356        dummy_inputs[graph_input.name] = data
357
358    all_inputs = [dummy_inputs for _ in range(samples)]
359    return all_inputs
360
361
362def run(args):
363    num_threads = args.thread_num if args.thread_num > 0 else psutil.cpu_count(logical=False)
364
365    # Set OMP environment variable before importing onnxruntime. Needed for cpu only, and no impact for onnxruntime-gpu package.
366    if "OMP_NUM_THREADS" not in os.environ:
367        os.environ["OMP_NUM_THREADS"] = str(num_threads)
368
369    from onnx import load  # noqa: PLC0415
370    from onnx_model import OnnxModel  # noqa: PLC0415
371
372    onnx_model = OnnxModel(load(args.model))
373
374    all_inputs = None
375    if args.dummy_inputs == "bert":
376        all_inputs = create_bert_inputs(
377            onnx_model,
378            args.batch_size,
379            args.sequence_length,
380            args.samples,
381            args.input_ids_name,
382            args.segment_ids_name,
383            args.input_mask_name,
384        )
385    elif args.dummy_inputs == "gpt2":
386        all_inputs = create_gpt2_inputs(
387            onnx_model,
388            args.batch_size,
389            args.sequence_length,
390            args.past_sequence_length,
391            args.samples,
392        )
393    elif args.dummy_inputs == "longformer":
394        all_inputs = create_longformer_inputs(
395            onnx_model,
396            args.batch_size,
397            args.sequence_length,
398            args.global_length,
399            args.samples,
400        )
401    else:  # default
402        all_inputs = create_dummy_inputs(onnx_model, args.batch_size, args.sequence_length, args.samples)
403
404    profile_file = run_profile(
405        args.model,
406        args.use_gpu,
407        args.provider,
408        args.basic_optimization,
409        args.thread_num,
410        all_inputs,
411    )
412
413    return profile_file
414
415
416if __name__ == "__main__":
417    arguments = parse_arguments()
418    print("Arguments", arguments)
419
420    from benchmark_helper import setup_logger
421
422    setup_logger(arguments.verbose)
423
424    if not arguments.input:
425        assert arguments.model, "requires either --model to run profiling or --input to read profiling results"
426        profile_file = run(arguments)
427    else:
428        profile_file = arguments.input
429    from profile_result_processor import process_results
430
431    results = process_results(profile_file, arguments)
432
433    for line in results:
434        print(line)
435