|
| 1 | +import logging |
| 2 | +import math |
| 3 | +from dataclasses import dataclass, field |
| 4 | +from typing import List, Tuple |
| 5 | + |
| 6 | +import torch |
| 7 | + |
| 8 | +logger = logging.getLogger(__name__) |
| 9 | + |
| 10 | + |
| 11 | +@dataclass |
| 12 | +class PerSubgraphData: |
| 13 | + """Class to track data on a per-subgraph level |
| 14 | +
|
| 15 | + Args: |
| 16 | + subgraph_name (str): Name of the subgraph in the GraphModule |
| 17 | + subgraph_op_count (int): Number of operations in the subgraph |
| 18 | + subgraph_input_shapes (List[Tuple[int, ...]]): Shapes of input Tensors of the subgraph |
| 19 | + subgraph_input_dtypes (List[torch.device]): Input data types of the subgraph |
| 20 | + subgraph_output_shapes (List[Tuple[int, ...]]): Shapes of output Tensors of the subgraph |
| 21 | + subgraph_output_dtypes (List[torch.device]): Output data types of the subgraph |
| 22 | + """ |
| 23 | + |
| 24 | + subgraph_name: str = "" |
| 25 | + subgraph_op_count: int = 0 |
| 26 | + subgraph_input_shapes: List[Tuple[int, ...]] = field(default_factory=list) |
| 27 | + subgraph_input_dtypes: List[torch.device] = field(default_factory=list) |
| 28 | + subgraph_output_shapes: List[Tuple[int, ...]] = field(default_factory=list) |
| 29 | + subgraph_output_dtypes: List[torch.device] = field(default_factory=list) |
| 30 | + |
| 31 | + |
| 32 | +@dataclass |
| 33 | +class DryRunTracker: |
| 34 | + """Class to track data on a graph-wide level |
| 35 | +
|
| 36 | + Args: |
| 37 | + total_ops_in_graph (int): Total number of operators in graph |
| 38 | + supported_ops_in_graph (int): Number of supported operators in graph |
| 39 | + graph_input_shapes (List[Tuple[int, ...]]): Shapes of input Tensors of the graph |
| 40 | + graph_input_dtypes (List[torch.device]): Input data types of the graph |
| 41 | + graph_output_shapes (List[Tuple[int, ...]]): Shapes of output Tensors of the graph |
| 42 | + graph_output_dtypes (List[torch.device]): Output data types of the graph |
| 43 | + per_subgraph_data (List[PerSubgraphData]): Per-subgraph data, see above class |
| 44 | + tensorrt_graph_count (int): Number of TensorRT engines to be generated |
| 45 | + truncated_long_and_double (bool): Whether truncate_long_and_double was enabled |
| 46 | + """ |
| 47 | + |
| 48 | + total_ops_in_graph: int = 0 |
| 49 | + supported_ops_in_graph: int = 0 |
| 50 | + graph_input_shapes: List[Tuple[int, ...]] = field(default_factory=list) |
| 51 | + graph_input_dtypes: List[torch.device] = field(default_factory=list) |
| 52 | + graph_output_shapes: List[Tuple[int, ...]] = field(default_factory=list) |
| 53 | + graph_output_dtypes: List[torch.device] = field(default_factory=list) |
| 54 | + per_subgraph_data: List[PerSubgraphData] = field(default_factory=list) |
| 55 | + tensorrt_graph_count: int = 0 |
| 56 | + truncated_long_and_double: bool = False |
| 57 | + |
| 58 | + |
| 59 | +def dryrun_stats_display(dryrun_tracker: DryRunTracker, dryrun_enabled: bool) -> None: |
| 60 | + """Displays statistics about the dryrun either to debug logs or info logs""" |
| 61 | + # If user specified "dryrun=True", print to info logs, else debug |
| 62 | + if dryrun_enabled: |
| 63 | + dryrun_logger = logger.info |
| 64 | + else: |
| 65 | + dryrun_logger = logger.debug |
| 66 | + |
| 67 | + formatted_stats = "\n" |
| 68 | + |
| 69 | + # Print overall stats about the graph, operator counts, etc. |
| 70 | + formatted_stats += "+" * 50 + " Dry-Run Results for Graph " + "+" * 50 + "\n" |
| 71 | + formatted_stats += ( |
| 72 | + f"The graph consists of {dryrun_tracker.total_ops_in_graph} Total Operators, " |
| 73 | + f"of which {dryrun_tracker.supported_ops_in_graph} operators are supported, " |
| 74 | + f"{round(dryrun_tracker.supported_ops_in_graph*100/dryrun_tracker.total_ops_in_graph, 2)}% coverage\n" |
| 75 | + ) |
| 76 | + formatted_stats += f"Long and double inputs were {'' if dryrun_tracker.truncated_long_and_double else 'not'} truncated (truncate_long_and_double={dryrun_tracker.truncated_long_and_double})\n" |
| 77 | + formatted_stats += ( |
| 78 | + f"{dryrun_tracker.tensorrt_graph_count} TRT Engine(s) were generated\n" |
| 79 | + ) |
| 80 | + |
| 81 | + assert len(dryrun_tracker.per_subgraph_data) == dryrun_tracker.tensorrt_graph_count |
| 82 | + |
| 83 | + # Print schematic of the graph structure, as in: |
| 84 | + # |
| 85 | + # Inputs: [Tensor: (1, 3, 224, 224)@float32] |
| 86 | + # ... |
| 87 | + # TRT Engine #1: _run_on_acc_0 |
| 88 | + # Engine Inputs: [Tensor: (1, 3, 224, 224)@float32] |
| 89 | + # Number of Operators in Engine: 1 |
| 90 | + # Engine Outputs: [Tensor: (1, 64, 112, 112)@float32] |
| 91 | + # ... |
| 92 | + # Outputs: [Tensor: (1, 1000)@float32] |
| 93 | + # |
| 94 | + formatted_stats += " " * 2 + "Graph Structure:\n\n" |
| 95 | + formatted_stats += ( |
| 96 | + " " * 3 |
| 97 | + + f"Inputs: [{input_formatter(dryrun_tracker.graph_input_shapes, dryrun_tracker.graph_input_dtypes)}]\n" |
| 98 | + ) |
| 99 | + |
| 100 | + for i, trt_subgraph_data in enumerate(dryrun_tracker.per_subgraph_data): |
| 101 | + assert len(trt_subgraph_data.subgraph_input_dtypes) == len( |
| 102 | + trt_subgraph_data.subgraph_input_shapes |
| 103 | + ) |
| 104 | + assert len(trt_subgraph_data.subgraph_output_dtypes) == len( |
| 105 | + trt_subgraph_data.subgraph_output_shapes |
| 106 | + ) |
| 107 | + formatted_stats += " " * 4 + "...\n" |
| 108 | + formatted_stats += ( |
| 109 | + " " * 4 + f"TRT Engine #{i+1}: {trt_subgraph_data.subgraph_name}\n" |
| 110 | + ) |
| 111 | + formatted_stats += ( |
| 112 | + " " * 5 |
| 113 | + + f"Engine Inputs: [{input_formatter(trt_subgraph_data.subgraph_input_shapes, trt_subgraph_data.subgraph_input_dtypes)}]\n" |
| 114 | + ) |
| 115 | + formatted_stats += ( |
| 116 | + " " * 5 |
| 117 | + + f"Number of Operators in Engine: {trt_subgraph_data.subgraph_op_count}\n" |
| 118 | + ) |
| 119 | + formatted_stats += ( |
| 120 | + " " * 5 |
| 121 | + + f"Engine Outputs: [{input_formatter(trt_subgraph_data.subgraph_output_shapes, trt_subgraph_data.subgraph_output_dtypes)}]\n" |
| 122 | + ) |
| 123 | + |
| 124 | + formatted_stats += " " * 4 + "...\n" |
| 125 | + formatted_stats += ( |
| 126 | + " " * 3 |
| 127 | + + f"Outputs: [{input_formatter(dryrun_tracker.graph_output_shapes, dryrun_tracker.graph_output_dtypes)}]\n" |
| 128 | + ) |
| 129 | + |
| 130 | + # Print aggregate statistics about the graph structure, including recommended "min_block_size" options |
| 131 | + if dryrun_tracker.tensorrt_graph_count > 0: |
| 132 | + min_ops_in_an_engine = min( |
| 133 | + trt_subgraph.subgraph_op_count |
| 134 | + for trt_subgraph in dryrun_tracker.per_subgraph_data |
| 135 | + ) |
| 136 | + avg_ops_per_engine = ( |
| 137 | + sum( |
| 138 | + trt_subgraph.subgraph_op_count |
| 139 | + for trt_subgraph in dryrun_tracker.per_subgraph_data |
| 140 | + ) |
| 141 | + / dryrun_tracker.tensorrt_graph_count |
| 142 | + ) |
| 143 | + avg_ops_per_engine = round(avg_ops_per_engine, 2) |
| 144 | + most_ops_in_an_engine = max( |
| 145 | + trt_subgraph.subgraph_op_count |
| 146 | + for trt_subgraph in dryrun_tracker.per_subgraph_data |
| 147 | + ) |
| 148 | + |
| 149 | + formatted_stats += "\n" + " " * 2 + "-" * 25 + " Aggregate Stats " + "-" * 25 |
| 150 | + formatted_stats += ( |
| 151 | + "\n\n" |
| 152 | + + " " * 3 |
| 153 | + + "Average Number of Operators per TRT Engine: " |
| 154 | + + f"{avg_ops_per_engine}" |
| 155 | + ) |
| 156 | + |
| 157 | + formatted_stats += ( |
| 158 | + "\n" |
| 159 | + + " " * 3 |
| 160 | + + "Most Operators in a TRT Engine: " |
| 161 | + + f"{most_ops_in_an_engine}" |
| 162 | + ) |
| 163 | + |
| 164 | + formatted_stats += "\n\n" + " " * 2 + "*" * 10 + " Recommendations " + "*" * 10 |
| 165 | + formatted_stats += ( |
| 166 | + "\n\n" |
| 167 | + + " " * 3 |
| 168 | + + "- For minimal graph segmentation, select min_block_size=" |
| 169 | + + f"{most_ops_in_an_engine} which would generate " |
| 170 | + + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= most_ops_in_an_engine])} TRT engines" |
| 171 | + ) |
| 172 | + if math.ceil(avg_ops_per_engine) != most_ops_in_an_engine: |
| 173 | + formatted_stats += ( |
| 174 | + "\n" |
| 175 | + + " " * 3 |
| 176 | + + "- For moderate graph segmentation, select min_block_size=" |
| 177 | + + f"{math.ceil(avg_ops_per_engine)} which would generate " |
| 178 | + + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= math.ceil(avg_ops_per_engine)])} TRT engines" |
| 179 | + ) |
| 180 | + |
| 181 | + formatted_stats += ( |
| 182 | + "\n" |
| 183 | + + " " * 3 |
| 184 | + + "- The current level of graph segmentation is equivalent to selecting min_block_size=" |
| 185 | + + f"{min_ops_in_an_engine} which generates " |
| 186 | + + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= min_ops_in_an_engine])} TRT engines" |
| 187 | + ) |
| 188 | + else: |
| 189 | + formatted_stats += ( |
| 190 | + "\n" |
| 191 | + + " " * 2 |
| 192 | + + "Aggregate stats not available since no TRT Engines were generated." |
| 193 | + ) |
| 194 | + |
| 195 | + dryrun_logger(formatted_stats) |
| 196 | + |
| 197 | + |
| 198 | +def input_formatter(shapes: List[Tuple[int, ...]], dtypes: List[torch.dtype]) -> str: |
| 199 | + """Format shapes and dtypes of input Tensors into a readable string""" |
| 200 | + formatted_str = ", " |
| 201 | + |
| 202 | + for shape, dtype in zip(shapes, dtypes): |
| 203 | + formatted_str += f"Tensor: {shape}@{str(dtype)[6:]}, " |
| 204 | + |
| 205 | + return formatted_str[2:-2] |
0 commit comments