from typing import Dict, Any, List import torch from collections import defaultdict from torch import nn import copy from ...sparsifier.utils import fqn_to_module, module_to_fqn import warnings __all__ = ['ActivationSparsifier'] class ActivationSparsifier: r""" The Activation sparsifier class aims to sparsify/prune activations in a neural network. The idea is to attach the sparsifier to a layer (or layers) and it zeroes out the activations based on the mask_fn (or sparsification function) input by the user. The mask_fn is applied once all the inputs are aggregated and reduced i.e. mask = mask_fn(reduce_fn(aggregate_fn(activations))) Note:: The sparsification mask is computed on the input **before it goes through the attached layer**. Args: model (nn.Module): The model whose layers will be sparsified. The layers that needs to be sparsified should be added separately using the register_layer() function aggregate_fn (Optional, Callable): default aggregate_fn that is used if not specified while registering the layer. specifies how inputs should be aggregated over time. The aggregate_fn should usually take 2 torch tensors and return the aggregated tensor. Example def add_agg_fn(tensor1, tensor2): return tensor1 + tensor2 reduce_fn (Optional, Callable): default reduce_fn that is used if not specified while registering the layer. reduce_fn will be called on the aggregated tensor i.e. the tensor obtained after calling agg_fn() on all inputs. Example def mean_reduce_fn(agg_tensor): return agg_tensor.mean(dim=0) mask_fn (Optional, Callable): default mask_fn that is used to create the sparsification mask using the tensor obtained after calling the reduce_fn(). This is used by default if a custom one is passed in the register_layer(). Note that the mask_fn() definition should contain the sparse arguments that is passed in sparse_config arguments. features (Optional, list): default selected features to sparsify. If this is non-empty, then the mask_fn will be applied for each feature of the input. For example, mask = [mask_fn(reduce_fn(aggregated_fn(input[feature])) for feature in features] feature_dim (Optional, int): default dimension of input features. Again, features along this dim will be chosen for sparsification. sparse_config (Dict): Default configuration for the mask_fn. This config will be passed with the mask_fn() Example: >>> # xdoctest: +SKIP >>> model = SomeModel() >>> act_sparsifier = ActivationSparsifier(...) # init activation sparsifier >>> # Initialize aggregate_fn >>> def agg_fn(x, y): >>> return x + y >>> >>> # Initialize reduce_fn >>> def reduce_fn(x): >>> return torch.mean(x, dim=0) >>> >>> # Initialize mask_fn >>> def mask_fn(data): >>> return torch.eye(data.shape).to(data.device) >>> >>> >>> act_sparsifier.register_layer(model.some_layer, aggregate_fn=agg_fn, reduce_fn=reduce_fn, mask_fn=mask_fn) >>> >>> # start training process >>> for _ in [...]: >>> # epoch starts >>> # model.forward(), compute_loss() and model.backwards() >>> # epoch ends >>> act_sparsifier.step() >>> # end training process >>> sparsifier.squash_mask() """ def __init__(self, model: nn.Module, aggregate_fn=None, reduce_fn=None, mask_fn=None, features=None, feature_dim=None, **sparse_config): self.model = model self.defaults: Dict[str, Any] = defaultdict() self.defaults['sparse_config'] = sparse_config # functions self.defaults['aggregate_fn'] = aggregate_fn self.defaults['reduce_fn'] = reduce_fn self.defaults['mask_fn'] = mask_fn # default feature and feature_dim self.defaults['features'] = features self.defaults['feature_dim'] = feature_dim self.data_groups: Dict[str, Dict] = defaultdict(dict) # contains all relevant info w.r.t each registered layer self.state: Dict[str, Any] = defaultdict(dict) # layer name -> mask @staticmethod def _safe_rail_checks(args): """Makes sure that some of the functions and attributes are not passed incorrectly """ # if features are not None, then feature_dim must not be None features, feature_dim = args['features'], args['feature_dim'] if features is not None: assert feature_dim is not None, "need feature dim to select features" # all the *_fns should be callable fn_keys = ['aggregate_fn', 'reduce_fn', 'mask_fn'] for key in fn_keys: fn = args[key] assert callable(fn), 'function should be callable' def _aggregate_hook(self, name): """Returns hook that computes aggregate of activations passing through. """ # gather some data feature_dim = self.data_groups[name]['feature_dim'] features = self.data_groups[name]['features'] agg_fn = self.data_groups[name]['aggregate_fn'] def hook(module, input) -> None: input_data = input[0] data = self.data_groups[name].get('data') # aggregated data if features is None: # no features associated, data should not be a list if data is None: data = torch.zeros_like(input_data) self.state[name]['mask'] = torch.ones_like(input_data) out_data = agg_fn(data, input_data) else: # data should be a list [aggregated over each feature only] if data is None: out_data = [0 for _ in range(0, len(features))] # create one incase of 1st forward self.state[name]['mask'] = [0 for _ in range(0, len(features))] else: out_data = data # a list # compute aggregate over each feature for feature_idx in range(len(features)): # each feature is either a list or scalar, convert it to torch tensor feature_tensor = torch.Tensor([features[feature_idx]]).long().to(input_data.device) data_feature = torch.index_select(input_data, feature_dim, feature_tensor) if data is None: curr_data = torch.zeros_like(data_feature) self.state[name]['mask'][feature_idx] = torch.ones_like(data_feature) else: curr_data = data[feature_idx] out_data[feature_idx] = agg_fn(curr_data, data_feature) self.data_groups[name]['data'] = out_data return hook def register_layer(self, layer: nn.Module, aggregate_fn=None, reduce_fn=None, mask_fn=None, features=None, feature_dim=None, **sparse_config): r""" Registers a layer for sparsification. The layer should be part of self.model. Specifically, registers a pre-forward hook to the layer. The hook will apply the aggregate_fn and store the aggregated activations that is input over each step. Note:: - There is no need to pass in the name of the layer as it is automatically computed as per the fqn convention. - All the functions (fn) passed as argument will be called at a dim, feature level. """ name = module_to_fqn(self.model, layer) assert name is not None, "layer not found in the model" # satisfy mypy if name in self.data_groups: # unregister layer if already present warnings.warn("layer already attached to the sparsifier, deregistering the layer and registering with new config") self.unregister_layer(name=name) local_args = copy.deepcopy(self.defaults) update_dict = { 'aggregate_fn': aggregate_fn, 'reduce_fn': reduce_fn, 'mask_fn': mask_fn, 'features': features, 'feature_dim': feature_dim, 'layer': layer } local_args.update((arg, val) for arg, val in update_dict.items() if val is not None) local_args['sparse_config'].update(sparse_config) self._safe_rail_checks(local_args) self.data_groups[name] = local_args agg_hook = layer.register_forward_pre_hook(self._aggregate_hook(name=name)) self.state[name]['mask'] = None # mask will be created when model forward is called. # attach agg hook self.data_groups[name]['hook'] = agg_hook # for serialization purposes, we know whether aggregate_hook is attached # or sparsify_hook() self.data_groups[name]['hook_state'] = "aggregate" # aggregate hook is attached def get_mask(self, name: str = None, layer: nn.Module = None): """ Returns mask associated to the layer. The mask is - a torch tensor is features for that layer is None. - a list of torch tensors for each feature, otherwise Note:: The shape of the mask is unknown until model.forward() is applied. Hence, if get_mask() is called before model.forward(), an error will be raised. """ assert name is not None or layer is not None, "Need at least name or layer obj to retrieve mask" if name is None: assert layer is not None name = module_to_fqn(self.model, layer) assert name is not None, "layer not found in the specified model" if name not in self.state: raise ValueError("Error: layer with the given name not found") mask = self.state[name].get('mask', None) if mask is None: raise ValueError("Error: shape unknown, call layer() routine at least once to infer mask") return mask def unregister_layer(self, name): """Detaches the sparsifier from the layer """ # detach any hooks attached self.data_groups[name]['hook'].remove() # pop from the state dict self.state.pop(name) # pop from the data groups self.data_groups.pop(name) def step(self): """Internally calls the update_mask() function for each layer """ with torch.no_grad(): for name, configs in self.data_groups.items(): data = configs['data'] self.update_mask(name, data, configs) self.data_groups[name].pop('data') # reset the accumulated data def update_mask(self, name, data, configs): """ Called for each registered layer and does the following- 1. apply reduce_fn on the aggregated activations 2. use mask_fn to compute the sparsification mask Note: the reduce_fn and mask_fn is called for each feature, dim over the data """ mask = self.get_mask(name) sparse_config = configs['sparse_config'] features = configs['features'] reduce_fn = configs['reduce_fn'] mask_fn = configs['mask_fn'] if features is None: data = reduce_fn(data) mask.data = mask_fn(data, **sparse_config) else: for feature_idx in range(len(features)): data_feature = reduce_fn(data[feature_idx]) mask[feature_idx].data = mask_fn(data_feature, **sparse_config) def _sparsify_hook(self, name): """Returns hook that applies sparsification mask to input entering the attached layer """ mask = self.get_mask(name) features = self.data_groups[name]['features'] feature_dim = self.data_groups[name]['feature_dim'] def hook(module, input): input_data = input[0] if features is None: # apply to all the features return input_data * mask else: # apply per feature, feature_dim for feature_idx in range(0, len(features)): feature = torch.Tensor([features[feature_idx]]).long().to(input_data.device) sparsified = torch.index_select(input_data, feature_dim, feature) * mask[feature_idx] input_data.index_copy_(feature_dim, feature, sparsified) return input_data return hook def squash_mask(self, attach_sparsify_hook=True, **kwargs): """ Unregisters aggreagate hook that was applied earlier and registers sparsification hooks if attach_sparsify_hook = True. """ for name, configs in self.data_groups.items(): # unhook agg hook configs['hook'].remove() configs.pop('hook') self.data_groups[name]['hook_state'] = "None" if attach_sparsify_hook: configs['hook'] = configs['layer'].register_forward_pre_hook(self._sparsify_hook(name)) configs['hook_state'] = "sparsify" # signals that sparsify hook is now attached def _get_serializable_data_groups(self): """Exclude hook and layer from the config keys before serializing TODO: Might have to treat functions (reduce_fn, mask_fn etc) in a different manner while serializing. For time-being, functions are treated the same way as other attributes """ data_groups: Dict[str, Any] = defaultdict() for name, config in self.data_groups.items(): new_config = {key: value for key, value in config.items() if key not in ['hook', 'layer']} data_groups[name] = new_config return data_groups def _convert_mask(self, states_dict, sparse_coo=True): r"""Converts the mask to sparse coo or dense depending on the `sparse_coo` argument. If `sparse_coo=True`, then the mask is stored as sparse coo else dense tensor """ states = copy.deepcopy(states_dict) for _, state in states.items(): if state['mask'] is not None: if isinstance(state['mask'], List): for idx in range(len(state['mask'])): if sparse_coo: state['mask'][idx] = state['mask'][idx].to_sparse_coo() else: state['mask'][idx] = state['mask'][idx].to_dense() else: if sparse_coo: state['mask'] = state['mask'].to_sparse_coo() else: state['mask'] = state['mask'].to_dense() return states def state_dict(self) -> Dict[str, Any]: r"""Returns the state of the sparsifier as a :class:`dict`. It contains: * state - contains name -> mask mapping. * data_groups - a dictionary containing all config information for each layer * defaults - the default config while creating the constructor """ data_groups = self._get_serializable_data_groups() state = self._convert_mask(self.state) return { 'state': state, 'data_groups': data_groups, 'defaults': self.defaults } def load_state_dict(self, state_dict: Dict[str, Any]) -> None: r"""The load_state_dict() restores the state of the sparsifier based on the state_dict Args: * state_dict - the dictionary that to which the current sparsifier needs to be restored to """ state = state_dict['state'] data_groups, defaults = state_dict['data_groups'], state_dict['defaults'] self.__set_state__({'state': state, 'data_groups': data_groups, 'defaults': defaults}) def __get_state__(self) -> Dict[str, Any]: data_groups = self._get_serializable_data_groups() state = self._convert_mask(self.state) return { 'defaults': self.defaults, 'state': state, 'data_groups': data_groups, } def __set_state__(self, state: Dict[str, Any]) -> None: state['state'] = self._convert_mask(state['state'], sparse_coo=False) # convert mask to dense tensor self.__dict__.update(state) # need to attach layer and hook info into the data_groups for name, config in self.data_groups.items(): # fetch layer layer = fqn_to_module(self.model, name) assert layer is not None # satisfy mypy # if agg_mode is True, then layer in aggregate mode if "hook_state" in config and config['hook_state'] == "aggregate": hook = layer.register_forward_pre_hook(self._aggregate_hook(name)) elif "hook_state" in config and config["hook_state"] == "sparsify": hook = layer.register_forward_pre_hook(self._sparsify_hook(name)) config['layer'] = layer config['hook'] = hook def __repr__(self): format_string = self.__class__.__name__ + ' (' for name, config in self.data_groups.items(): format_string += '\n' format_string += '\tData Group\n' format_string += f'\t name: {name}\n' for key in sorted(config.keys()): if key in ['data', 'hook', 'reduce_fn', 'mask_fn', 'aggregate_fn']: continue format_string += f'\t {key}: {config[key]}\n' format_string += ')' return format_string