# mypy: allow-untyped-defs from collections import OrderedDict, OrderedDict as OrdDict from typing import Any import torch # try to import tablate got_tabulate = True try: from tabulate import tabulate except ImportError: got_tabulate = False # var to see if we could import matplotlib got_matplotlib = True try: import matplotlib.pyplot as plt except ImportError: got_matplotlib = False class ModelReportVisualizer: r""" The ModelReportVisualizer class aims to provide users a way to visualize some of the statistics that were generated by the ModelReport API. However, at a higher level, the class aims to provide some level of visualization of statistics to PyTorch in order to make it easier to parse data and diagnose any potential issues with data or a specific model. With respect to the visualizations, the ModelReportVisualizer class currently supports several methods of visualizing data. Supported Visualization Methods Include: - Table format - Plot format (line graph) - Histogram format For all of the existing visualization methods, there is the option to filter data based on: - A module fqn prefix - Feature [required for the plot and histogram] * :attr:`generated_reports` The reports generated by the ModelReport class in the structure below Ensure sure that features that are the same across different report contain the same name Ensure that objects representing the same features are the same type / dimension (where applicable) Note: Currently, the ModelReportVisualizer class supports visualization of data generated by the ModelReport class. However, this structure is extensible and should allow the visualization of other information as long as the information is structured in the following general format: Report Structure -- module_fqn [module with attached detectors] | -- feature keys [not every detector extracts same information] [same collected info has same keys, unless can be specific to detector] The goal behind the class is that the generated visualizations can be used in conjunction with the generated report for people to get a better understanding of issues and what the fix might be. It is also just to provide a good visualization platform, since it might be hard to parse through the ModelReport returned dictionary as that grows in size. General Use Flow Expected 1.) Initialize ModelReport object with reports of interest by passing in initialized detector objects 2.) Prepare your model with prepare_fx 3.) Call model_report.prepare_detailed_calibration on your model to add relevant observers 4.) Calibrate your model with data 5.) Call model_report.generate_report on your model to generate report and optionally remove added observers 6.) Use output of model_report.generate_report to initialize ModelReportVisualizer instance 7.) Use instance to view different views of data as desired, applying filters as needed 8.) Either see the super detailed information or just the actual printed or shown table / plot / histogram """ # keys for table dict TABLE_TENSOR_KEY = "tensor_level_info" TABLE_CHANNEL_KEY = "channel_level_info" # Constants for header vals NUM_NON_FEATURE_TENSOR_HEADERS = 2 NUM_NON_FEATURE_CHANNEL_HEADERS = 3 # Constants for row index in header CHANNEL_NUM_INDEX = 2 def __init__(self, generated_reports: OrderedDict[str, Any]): r""" Initializes the ModelReportVisualizer instance with the necessary reports. Args: generated_reports (Dict[str, Any]): The reports generated by the ModelReport class can also be a dictionary generated in another manner, as long as format is same """ self.generated_reports = generated_reports def get_all_unique_module_fqns(self) -> set[str]: r""" The purpose of this method is to provide a user the set of all module_fqns so that if they wish to use some of the filtering capabilities of the ModelReportVisualizer class, they don't need to manually parse the generated_reports dictionary to get this information. Returns all the unique module fqns present in the reports the ModelReportVisualizer instance was initialized with. """ # returns the keys of the ordered dict return set(self.generated_reports.keys()) def get_all_unique_feature_names( self, plottable_features_only: bool = True ) -> set[str]: r""" The purpose of this method is to provide a user the set of all feature names so that if they wish to use the filtering capabilities of the generate_table_view(), or use either of the generate_plot_view() or generate_histogram_view(), they don't need to manually parse the generated_reports dictionary to get this information. Args: plottable_features_only (bool): True if the user is only looking for plottable features, False otherwise plottable features are those that are tensor values Default: True (only return those feature names that are plottable) Returns all the unique module fqns present in the reports the ModelReportVisualizer instance was initialized with. """ unique_feature_names = set() for module_fqn in self.generated_reports: # get dict of the features feature_dict: dict[str, Any] = self.generated_reports[module_fqn] # loop through features for feature_name in feature_dict: # if we need plottable, ensure type of val is tensor if ( not plottable_features_only or type(feature_dict[feature_name]) is torch.Tensor ): unique_feature_names.add(feature_name) # return our compiled set of unique feature names return unique_feature_names def _get_filtered_data( self, feature_filter: str, module_fqn_filter: str ) -> OrderedDict[str, Any]: r""" Filters the data and returns it in the same ordered dictionary format so the relevant views can be displayed. Args: feature_filter (str): The feature filter, if we want to filter the set of data to only include a certain set of features that include feature_filter If feature = "", then we do not filter based on any features module_fqn_filter (str): The filter on prefix for the module fqn. All modules that have fqn with this prefix will be included If module_fqn_filter = "" we do not filter based on module fqn, and include all modules First, the data is filtered based on module_fqn, and then filtered based on feature Returns an OrderedDict (sorted in order of model) mapping: module_fqns -> feature_names -> values """ # create return dict filtered_dict: OrderedDict[str, Any] = OrdDict() for module_fqn in self.generated_reports: # first filter based on module if module_fqn_filter == "" or module_fqn_filter in module_fqn: # create entry for module and loop through features filtered_dict[module_fqn] = {} module_reports = self.generated_reports[module_fqn] for feature_name in module_reports: # check if filtering on features and do so if desired if feature_filter == "" or feature_filter in feature_name: filtered_dict[module_fqn][feature_name] = module_reports[ feature_name ] # we have populated the filtered dict, and must return it return filtered_dict def _generate_tensor_table( self, filtered_data: OrderedDict[str, dict[str, Any]], tensor_features: list[str], ) -> tuple[list, list]: r""" Takes in the filtered data and features list and generates the tensor headers and table Currently meant to generate the headers and table for both the tensor information. Args: filtered_data (OrderedDict[str, Dict[str, Any]]): An OrderedDict (sorted in order of model) mapping: module_fqns -> feature_names -> values tensor_features (List[str]): A list of the tensor level features Returns a tuple with: A list of the headers of the tensor table A list of lists containing the table information row by row The 0th index row will contain the headers of the columns The rest of the rows will contain data """ # now we compose the tensor information table tensor_table: list[list[Any]] = [] tensor_headers: list[str] = [] # append the table row to the table only if we have features if len(tensor_features) > 0: # now we add all the data for index, module_fqn in enumerate(filtered_data): # we make a new row for the tensor table tensor_table_row = [index, module_fqn] for feature in tensor_features: # we iterate in same order of added features if feature in filtered_data[module_fqn]: # add value if applicable to module feature_val = filtered_data[module_fqn][feature] else: # add that it is not applicable feature_val = "Not Applicable" # if it's a tensor we want to extract val if isinstance(feature_val, torch.Tensor): feature_val = feature_val.item() # we add to our list of values # pyrefly: ignore [bad-argument-type] tensor_table_row.append(feature_val) tensor_table.append(tensor_table_row) # add row of headers of we actually have something, otherwise just empty if len(tensor_table) != 0: tensor_headers = ["idx", "layer_fqn"] + tensor_features return (tensor_headers, tensor_table) def _generate_channels_table( self, filtered_data: OrderedDict[str, Any], channel_features: list[str], num_channels: int, ) -> tuple[list, list]: r""" Takes in the filtered data and features list and generates the channels headers and table Currently meant to generate the headers and table for both the channels information. Args: filtered_data (OrderedDict[str, Any]): An OrderedDict (sorted in order of model) mapping: module_fqns -> feature_names -> values channel_features (List[str]): A list of the channel level features num_channels (int): Number of channels in the channel data Returns a tuple with: A list of the headers of the channel table A list of lists containing the table information row by row The 0th index row will contain the headers of the columns The rest of the rows will contain data """ # now we compose the table for the channel information table channel_table: list[list[Any]] = [] channel_headers: list[str] = [] # counter to keep track of number of entries in channel_table_entry_counter: int = 0 if len(channel_features) > 0: # now we add all channel data for module_fqn in filtered_data: # we iterate over all channels for channel in range(num_channels): # we make a new row for the channel new_channel_row = [channel_table_entry_counter, module_fqn, channel] for feature in channel_features: if feature in filtered_data[module_fqn]: # add value if applicable to module feature_val = filtered_data[module_fqn][feature][channel] else: # add that it is not applicable feature_val = "Not Applicable" # if it's a tensor we want to extract val if type(feature_val) is torch.Tensor: feature_val = feature_val.item() # add value to channel specific row # pyrefly: ignore [bad-argument-type] new_channel_row.append(feature_val) # add to table and increment row index counter channel_table.append(new_channel_row) channel_table_entry_counter += 1 # add row of headers of we actually have something, otherwise just empty if len(channel_table) != 0: channel_headers = ["idx", "layer_fqn", "channel"] + channel_features return (channel_headers, channel_table) def generate_filtered_tables( self, feature_filter: str = "", module_fqn_filter: str = "" ) -> dict[str, tuple[list, list]]: r""" Takes in optional filter values and generates two tables with desired information. The generated tables are presented in both a list-of-lists format The reason for the two tables are that they handle different things: 1.) the first table handles all tensor level information 2.) the second table handles and displays all channel based information The reasoning for this is that having all the info in one table can make it ambiguous which collected statistics are global, and which are actually per-channel, so it's better to split it up into two tables. This also makes the information much easier to digest given the plethora of statistics collected Tensor table columns: idx layer_fqn feature_1 feature_2 feature_3 .... feature_n ---- --------- --------- --------- --------- --------- Per-Channel table columns: idx layer_fqn channel feature_1 feature_2 feature_3 .... feature_n ---- --------- ------- --------- --------- --------- --------- Args: feature_filter (str, optional): Filters the features presented to only those that contain this filter substring Default = "", results in all the features being printed module_fqn_filter (str, optional): Only includes modules that contains this string Default = "", results in all the modules in the reports to be visible in the table Returns a dictionary with two keys: (Dict[str, Tuple[List, List]]) A dict containing two keys: "tensor_level_info", "channel_level_info" Each key maps to a tuple with: A list of the headers of each table A list of lists containing the table information row by row The 0th index row will contain the headers of the columns The rest of the rows will contain data Example Use: >>> # xdoctest: +SKIP("undefined variables") >>> mod_report_visualizer.generate_filtered_tables( ... feature_filter="per_channel_min", module_fqn_filter="block1" ... ) # generates table with per_channel_min info for all modules in block 1 of the model """ # first get the filtered data filtered_data: OrderedDict[str, Any] = self._get_filtered_data( feature_filter, module_fqn_filter ) # now we split into tensor and per-channel data tensor_features: set[str] = set() channel_features: set[str] = set() # keep track of the number of channels we have num_channels: int = 0 for module_fqn in filtered_data: for feature_name in filtered_data[module_fqn]: # get the data for that specific feature feature_data = filtered_data[module_fqn][feature_name] # check if not zero dim tensor is_tensor: bool = isinstance(feature_data, torch.Tensor) is_not_zero_dim: bool = is_tensor and len(feature_data.shape) != 0 if is_not_zero_dim or isinstance(feature_data, list): # works means per channel channel_features.add(feature_name) num_channels = len(feature_data) else: # means is per-tensor tensor_features.add(feature_name) # we make them lists for iteration purposes tensor_features_list: list[str] = sorted(tensor_features) channel_features_list: list[str] = sorted(channel_features) # get the tensor info tensor_headers, tensor_table = self._generate_tensor_table( filtered_data, tensor_features_list ) # get the channel info channel_headers, channel_table = self._generate_channels_table( filtered_data, channel_features_list, num_channels ) # let's now create the dictionary to return table_dict = { self.TABLE_TENSOR_KEY: (tensor_headers, tensor_table), self.TABLE_CHANNEL_KEY: (channel_headers, channel_table), } # return the two tables return table_dict def generate_table_visualization( self, feature_filter: str = "", module_fqn_filter: str = "" ): r""" Takes in optional filter values and prints out formatted tables of the information. The reason for the two tables printed out instead of one large one are that they handle different things: 1.) the first table handles all tensor level information 2.) the second table handles and displays all channel based information The reasoning for this is that having all the info in one table can make it ambiguous which collected statistics are global, and which are actually per-channel, so it's better to split it up into two tables. This also makes the information much easier to digest given the plethora of statistics collected Tensor table columns: idx layer_fqn feature_1 feature_2 feature_3 .... feature_n ---- --------- --------- --------- --------- --------- Per-Channel table columns: idx layer_fqn channel feature_1 feature_2 feature_3 .... feature_n ---- --------- ------- --------- --------- --------- --------- Args: feature_filter (str, optional): Filters the features presented to only those that contain this filter substring Default = "", results in all the features being printed module_fqn_filter (str, optional): Only includes modules that contains this string Default = "", results in all the modules in the reports to be visible in the table Example Use: >>> # xdoctest: +SKIP("undefined variables") >>> mod_report_visualizer.generate_table_visualization( ... feature_filter="per_channel_min", module_fqn_filter="block1" ... ) >>> # prints out neatly formatted table with per_channel_min info >>> # for all modules in block 1 of the model """ # see if we got tabulate if not got_tabulate: print("Make sure to install tabulate and try again.") return None # get the table dict and the specific tables of interest table_dict = self.generate_filtered_tables(feature_filter, module_fqn_filter) tensor_headers, tensor_table = table_dict[self.TABLE_TENSOR_KEY] channel_headers, channel_table = table_dict[self.TABLE_CHANNEL_KEY] # get the table string and print it out # now we have populated the tables for each one # let's create the strings to be returned table_str = "" # the tables will have some headers columns that are non-feature # ex. table index, module name, channel index, etc. # we want to look at header columns for features, that come after those headers if len(tensor_headers) > self.NUM_NON_FEATURE_TENSOR_HEADERS: # if we have at least one tensor level feature to be added we add tensor table table_str += "Tensor Level Information \n" table_str += tabulate(tensor_table, headers=tensor_headers) if len(channel_headers) > self.NUM_NON_FEATURE_CHANNEL_HEADERS: # if we have at least one channel level feature to be added we add tensor table table_str += "\n\n Channel Level Information \n" table_str += tabulate(channel_table, headers=channel_headers) # if no features at all, let user know if table_str == "": table_str = "No data points to generate table with." print(table_str) def _get_plottable_data( self, feature_filter: str, module_fqn_filter: str ) -> tuple[list, list[list], bool]: r""" Takes in the feature filters and module filters and outputs the x and y data for plotting Args: feature_filter (str): Filters the features presented to only those that contain this filter substring module_fqn_filter (str): Only includes modules that contains this string Returns a tuple of three elements The first is a list containing relevant x-axis data The second is a list containing the corresponding y-axis data If the data is per channel """ # get the table dict and the specific tables of interest table_dict = self.generate_filtered_tables(feature_filter, module_fqn_filter) tensor_headers, tensor_table = table_dict[self.TABLE_TENSOR_KEY] channel_headers, channel_table = table_dict[self.TABLE_CHANNEL_KEY] # make sure it is only 1 feature that is being plotted # get the number of features in each of these tensor_info_features_count = ( len(tensor_headers) - ModelReportVisualizer.NUM_NON_FEATURE_TENSOR_HEADERS ) channel_info_features_count = ( len(channel_headers) - ModelReportVisualizer.NUM_NON_FEATURE_CHANNEL_HEADERS ) # see if valid tensor or channel plot is_valid_per_tensor_plot: bool = tensor_info_features_count == 1 is_valid_per_channel_plot: bool = channel_info_features_count == 1 # offset should either be one of tensor or channel table or neither feature_column_offset = ModelReportVisualizer.NUM_NON_FEATURE_TENSOR_HEADERS table = tensor_table # if a per_channel plot, we have different offset and table if is_valid_per_channel_plot: feature_column_offset = ( ModelReportVisualizer.NUM_NON_FEATURE_CHANNEL_HEADERS ) table = channel_table x_data: list = [] y_data: list[list] = [] # the feature will either be a tensor feature or channel feature if is_valid_per_tensor_plot: for table_row_num, row in enumerate(table): # get x_value to append x_val_to_append = table_row_num # the index of the feature will the 0 + num non feature columns tensor_feature_index = feature_column_offset row_value = row[tensor_feature_index] if type(row_value) is not str: x_data.append(x_val_to_append) y_data.append(row_value) elif is_valid_per_channel_plot: # gather the x_data and multiple y_data # calculate the number of channels num_channels: int = max(row[self.CHANNEL_NUM_INDEX] for row in table) + 1 # separate data list per channel y_data.extend([] for _ in range(num_channels)) for table_row_num, row in enumerate(table): # get x_value to append x_val_to_append = table_row_num current_channel = row[ self.CHANNEL_NUM_INDEX ] # initially chose current channel new_module_index: int = table_row_num // num_channels x_val_to_append = new_module_index # the index of the feature will the 0 + num non feature columns tensor_feature_index = feature_column_offset row_value = row[tensor_feature_index] if type(row_value) is not str: # only append if new index we are appending if len(x_data) == 0 or x_data[-1] != x_val_to_append: x_data.append(x_val_to_append) # append value for that channel y_data[current_channel].append(row_value) else: # more than one feature was chosen error_str = "Make sure to pick only a single feature with your filter to plot a graph." error_str += " We recommend calling get_all_unique_feature_names() to find unique feature names." error_str += " Pick one of those features to plot." raise ValueError(error_str) # return x, y values, and if data is per-channel return (x_data, y_data, is_valid_per_channel_plot) def generate_plot_visualization( self, feature_filter: str, module_fqn_filter: str = "" ): r""" Takes in a feature and optional module_filter and plots of the desired data. For per channel features, it averages the value across the channels and plots a point per module. The reason for this is that for models with hundreds of channels, it can be hard to differentiate one channel line from another, and so the point of generating a single average point per module is to give a sense of general trends that encourage further deep dives. Note: Only features in the report that have tensor value data are plottable by this class When the tensor information is plotted, it will plot: idx as the x val, feature value as the y_val When the channel information is plotted, it will plot: the first idx of each module as the x val, feature value as the y_val [for each channel] The reason for this is that we want to be able to compare values across the channels for same layer, and it will be hard if values are staggered by idx This means each module is represented by only 1 x value Args: feature_filter (str): Filters the features presented to only those that contain this filter substring module_fqn_filter (str, optional): Only includes modules that contains this string Default = "", results in all the modules in the reports to be visible in the table Example Use: >>> # xdoctest: +SKIP("undefined variables") >>> mod_report_visualizer.generate_plot_visualization( ... feature_filter="per_channel_min", module_fqn_filter="block1" ... ) >>> # outputs line plot of per_channel_min information for all >>> # modules in block1 of model each channel gets it's own line, >>> # and it's plotted across the in-order modules on the x-axis """ # checks if we have matplotlib and let's user know to install it if don't if not got_matplotlib: print("make sure to install matplotlib and try again.") return None # get the x and y data and if per channel x_data, y_data, data_per_channel = self._get_plottable_data( feature_filter, module_fqn_filter ) # plot based on whether data is per channel or not ax = plt.subplot() ax.set_ylabel(feature_filter) ax.set_title(feature_filter + " Plot") plt.xticks(x_data) # only show ticks for actual points if data_per_channel: ax.set_xlabel("First idx of module") # set the legend as well # plot a single line that is average of the channel values num_modules = len( y_data[0] ) # all y_data have same length, so get num modules num_channels = len( y_data ) # we want num channels to be able to calculate average later avg_vals = [ sum(y_data[:][index]) / num_channels for index in range(num_modules) ] # plot the three things we measured ax.plot( x_data, avg_vals, label=f"Average Value Across {num_channels} Channels" ) ax.legend(loc="upper right") else: ax.set_xlabel("idx") ax.plot(x_data, y_data) # actually show the plot plt.show() def generate_histogram_visualization( self, feature_filter: str, module_fqn_filter: str = "", num_bins: int = 10 ): r""" Takes in a feature and optional module_filter and plots the histogram of desired data. Note: Only features in the report that have tensor value data can be viewed as a histogram If you want to plot a histogram from all the channel values of a specific feature for a specific model, make sure to specify both the model and the feature properly in the filters and you should be able to see a distribution of the channel data Args: feature_filter (str, optional): Filters the features presented to only those that contain this filter substring Default = "", results in all the features being printed module_fqn_filter (str, optional): Only includes modules that contains this string Default = "", results in all the modules in the reports to be visible in the table num_bins (int, optional): The number of bins to create the histogram with Default = 10, the values will be split into 10 equal sized bins Example Use: >>> # xdoctest: +SKIP >>> mod_report_visualizer.generategenerate_histogram_visualization_plot_visualization( ... feature_filter="per_channel_min", module_fqn_filter="block1" ... ) # outputs histogram of per_channel_min information for all modules in block1 of model information is gathered across all channels for all modules in block 1 for the per_channel_min and is displayed in a histogram of equally sized bins """ # checks if we have matplotlib and let's user know to install it if don't if not got_matplotlib: print("make sure to install matplotlib and try again.") return None # get the x and y data and if per channel _x_data, y_data, data_per_channel = self._get_plottable_data( feature_filter, module_fqn_filter ) # for histogram, we just care about plotting the y data # plot based on whether data is per channel or not ax = plt.subplot() ax.set_xlabel(feature_filter) ax.set_ylabel("Frequency") ax.set_title(feature_filter + " Histogram") if data_per_channel: # set the legend as well # combine all the data all_data = [] for channel_info in y_data: all_data.extend(channel_info) _val, bins, _ = plt.hist( all_data, bins=num_bins, stacked=True, rwidth=0.8, ) plt.xticks(bins) else: _val, bins, _ = plt.hist( y_data, bins=num_bins, stacked=False, rwidth=0.8, ) plt.xticks(bins) plt.show()