import collections import typing from abc import ABC, abstractmethod from inspect import getsource from pprint import pformat import torch import torch.distributed as dist import torch.nn as nn from torch.nn.parallel.data_parallel import DataParallel from torch.nn.parallel.distributed import DistributedDataParallel from tinynn.graph.modifier import is_dw_conv from tinynn.graph.tracer import TraceGraph, trace from tinynn.util.train_util import DLContext, get_module_device from tinynn.util.util import conditional, get_actual_type, get_logger try: import ruamel_yaml as yaml except ModuleNotFoundError: import ruamel.yaml as yaml NEW_YAML_FLAG = "error_deprecation" in getsource(yaml.load) log = get_logger(__name__) class BasePruner(ABC): required_params = () required_context_params = () default_values = {} context_from_params_dict = {} condition_dict = {} model: nn.Module dummy_input: torch.Tensor config: typing.Union[dict, collections.OrderedDict] graph: TraceGraph context: DLContext def __init__(self, model, dummy_input, config): self.model = model self.config = config self.dummy_input = dummy_input self.graph = self.trace() @abstractmethod def prune(self): """The main function for pruning""" pass @abstractmethod def register_mask(self): """Computes the mask for the parameters in the model and register them through the maskers""" pass @abstractmethod def apply_mask(self): """Applies the masks for the parameters and updates the shape and properties of the tensors and modules""" pass def parse_config(self): """Parses the config and init the parameters of the pruner""" if isinstance(self.config, str): self.config = self.load_config(self.config) if not isinstance(self.config, dict): raise Exception('The `config` argument requires a parsed json object (e.g. dict or OrderedDict)') missing_params = set(self.required_params) - set(self.config) if len(missing_params) != 0: missing_params_str = ', '.join(missing_params) raise Exception(f'Missing param {missing_params_str} for {type(self).__name__}') for param_key, default_value in self.default_values.items(): if param_key not in self.config: self.config[param_key] = default_value type_dict = {} for cls in reversed(type(self).__mro__): if hasattr(cls, '__annotations__') and cls != BasePruner: type_dict.update(cls.__annotations__) for param_key, param_type in type_dict.items(): if param_key not in self.config: continue type_expected = False param_types = get_actual_type(param_type) for type_cand in param_types: if isinstance(self.config[param_key], type_cand): type_expected = True break if not type_expected: raise Exception(f'The type of `{param_key}` in {type(self).__name__} should be {param_type}') for param_key, predicate in self.condition_dict.items(): if predicate(self.config[param_key]) is False: raise Exception(f'The value of `{param_key}` doesn\'t meet the requirement: {getsource(predicate)}') def parse_context(self, context: DLContext): """Parses the context and copy the needed items to the pruner""" for context_key, param_keys in self.context_from_params_dict.items(): for param_key in param_keys: if param_key in self.config: setattr(context, context_key, self.config[param_key]) break filtered_context = dict(filter(lambda x: x[1], context.__dict__.items())) missing_context_items = list(set(self.required_context_params) - set(filtered_context)) if len(missing_context_items) != 0: missing_context_items_str = ', '.join(missing_context_items) raise Exception(f'Missing context items {missing_context_items_str} for {type(self).__name__}') self.context = context def summary(self): """Dumps the parameters and possibly the related context items of the pruner""" if len(self.required_params) > 0: log.info('-' * 80) log.info(f'params ({type(self).__name__}):') for k, v in self.__dict__.items(): if k in self.required_params or k in self.default_values: log.info(f'{k}: {pformat(v)}') if len(self.required_context_params) > 0: log.info('-' * 80) log.info(f'context ({type(self).__name__}):') log.info('\n'.join((f'{k}: {pformat(v)}' for k, v in self.context.__dict__.items()))) @classmethod def load_config(cls, path: str) -> dict: """Loads the configuration file and returns it as a dictionary""" with open(path, 'r') as f: if NEW_YAML_FLAG: yaml_ = yaml.YAML(typ='rt') config = yaml_.load(f) else: config = yaml.load(f, Loader=yaml.RoundTripLoader) return config @conditional(lambda: not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0) def generate_config(self, path: str, config: dict = None) -> None: """Generates a new copy the updated configuration with the given path""" if config is None: config = self.config with open(path, 'w') as f: if NEW_YAML_FLAG: yaml_ = yaml.YAML(typ='rt') yaml_.default_flow_style = False yaml_.dump(config, f) else: yaml.dump(config, f, default_flow_style=False, Dumper=yaml.RoundTripDumper) def trace(self) -> TraceGraph: with torch.no_grad(): if isinstance(self.model, DataParallel) or isinstance(self.model, DistributedDataParallel): model = self.model.module else: model = self.model old_device = get_module_device(model) model.cpu() graph = trace(model, self.dummy_input) if old_device is not None: model.to(device=old_device) return graph def reset(self): """Regenerate the TraceGraph when it is invalidated""" self.graph = self.trace() def calc_flops(self) -> int: """Calculate the flops of the given model""" # If graph is invalidated, then we need to regenerate the graph if not self.graph.inited: self.reset() graph: TraceGraph = self.graph total_ops = 0 for node in graph.forward_nodes: m = node.module remove_in_channel_count = 0 remove_out_channel_count = 0 if hasattr(m, 'masker'): if m.masker.in_remove_idx is not None: remove_in_channel_count = len(m.masker.in_remove_idx) if m.masker.ot_remove_idx is not None: remove_out_channel_count = len(m.masker.ot_remove_idx) if type(m) in (nn.Conv2d, nn.ConvTranspose2d): kernel_ops = torch.zeros(m.weight.size()[2:]).numel() bias_ops = 1 if m.bias is not None else 0 in_channels = m.in_channels - remove_in_channel_count out_channels = m.out_channels - remove_out_channel_count out_elements = node.next_tensors[0].nelement() // m.out_channels * out_channels if is_dw_conv(m): groups = in_channels else: groups = m.groups total_ops += out_elements * (in_channels // groups * kernel_ops + bias_ops) elif type(m) in (nn.BatchNorm2d,): channels = m.num_features - remove_in_channel_count nelements = node.prev_tensors[0].numel() // m.num_features * channels total_ops += 2 * nelements elif type(m) in (nn.AvgPool2d, nn.AdaptiveAvgPool2d): channels = node.prev_tensors[0].size(1) - remove_in_channel_count kernel_ops = 1 num_elements = node.prev_tensors[0].numel() // node.prev_tensors[0].size(1) * channels total_ops += kernel_ops * num_elements elif type(m) in (nn.ReLU,): channels = node.prev_tensors[0].size(1) - remove_in_channel_count kernel_ops = 1 num_elements = node.prev_tensors[0].numel() // node.prev_tensors[0].size(1) * channels total_ops += kernel_ops * num_elements elif type(m) in (nn.Linear,): in_channels = m.in_features - remove_in_channel_count out_channels = m.out_features - remove_out_channel_count total_mul = in_channels num_elements = node.next_tensors[0].numel() // m.out_features * out_channels total_ops += total_mul * num_elements return total_ops