1# mypy: allow-untyped-defs 2import torch 3from torch import nn 4from torch.nn.utils.parametrize import is_parametrized 5 6 7def module_contains_param(module, parametrization): 8 if is_parametrized(module): 9 # see if any of the module tensors have a parametriztion attached that matches the one passed in 10 return any( 11 any(isinstance(param, parametrization) for param in param_list) 12 for key, param_list in module.parametrizations.items() 13 ) 14 return False 15 16 17# Structured Pruning Parameterizations 18class FakeStructuredSparsity(nn.Module): 19 r""" 20 Parametrization for Structured Pruning. Like FakeSparsity, this should be attached to 21 the 'weight' or any other parameter that requires a mask. 22 23 Instead of an element-wise bool mask, this parameterization uses a row-wise bool mask. 24 """ 25 26 def __init__(self, mask): 27 super().__init__() 28 self.register_buffer("mask", mask) 29 30 def forward(self, x): 31 assert isinstance(self.mask, torch.Tensor) 32 assert self.mask.shape[0] == x.shape[0] 33 shape = [1] * len(x.shape) 34 shape[0] = -1 35 return self.mask.reshape(shape) * x 36 37 def state_dict(self, *args, **kwargs): 38 # avoid double saving masks 39 return {} 40 41 42class BiasHook: 43 def __init__(self, parametrization, prune_bias): 44 self.param = parametrization 45 self.prune_bias = prune_bias 46 47 def __call__(self, module, input, output): 48 if getattr(module, "_bias", None) is not None: 49 bias = module._bias.data 50 if self.prune_bias: 51 bias[~self.param.mask] = 0 52 53 # reshape bias to broadcast over output dimensions 54 idx = [1] * len(output.shape) 55 idx[1] = -1 56 bias = bias.reshape(idx) 57 58 output += bias 59 return output 60