from distutils.version import LooseVersion import numbers from typing import Optional, Tuple, Union import warnings import torch from torch import Tensor from tinynn.util.train_util import get_logger log = get_logger(__name__, 'WARNING') if LooseVersion(torch.__version__) >= '1.13.0': @classmethod def from_float(cls, other, qconfig=None): assert isinstance(other, cls._FLOAT_MODULE) assert hasattr(other, 'qconfig') or qconfig observed = cls( other.input_size, other.hidden_size, other.num_layers, other.bias, other.batch_first, other.dropout, other.bidirectional, ) observed.qconfig = getattr(other, 'qconfig', qconfig) for idx in range(other.num_layers): observed.layers[idx] = _GRULayer.from_float(other, idx, qconfig, batch_first=False) observed.train() observed = torch.ao.quantization.prepare_qat(observed, inplace=True) return observed class GRUCell(torch.nn.Module): r"""A quantizable gated recurrent unit (GRU) cell. For the description and the argument types, please, refer to :class:`~torch.nn.GRUCell` """ _FLOAT_MODULE = torch.nn.GRUCell def __init__(self, input_dim: int, hidden_dim: int, bias: bool = True, device=None, dtype=None) -> None: factory_kwargs = {'device': device, 'dtype': dtype} super().__init__() self.input_size = input_dim self.hidden_size = hidden_dim self.bias = bias self.igates = torch.nn.Linear(input_dim, 3 * hidden_dim, bias=bias, **factory_kwargs) self.hgates = torch.nn.Linear(hidden_dim, 3 * hidden_dim, bias=bias, **factory_kwargs) self.add1 = torch.ao.nn.quantized.FloatFunctional() self.add2 = torch.ao.nn.quantized.FloatFunctional() self.add3 = torch.ao.nn.quantized.FloatFunctional() self.add4 = torch.ao.nn.quantized.FloatFunctional() self.sub1 = torch.ao.nn.quantized.FloatFunctional() self.mul1 = torch.ao.nn.quantized.FloatFunctional() self.mul2 = torch.ao.nn.quantized.FloatFunctional() self.mul3 = torch.ao.nn.quantized.FloatFunctional() self.mul4 = torch.ao.nn.quantized.FloatFunctional() self.act1 = torch.nn.Sigmoid() self.act2 = torch.nn.Tanh() self.hidden_state_dtype: torch.dtype = torch.quint8 def forward(self, x: Tensor, hidden: Optional[Tensor] = None) -> Union[Tuple[Tensor, Tensor], Tensor]: result = [] if hidden is None or hidden[0] is None: hidden = self.initialize_hidden(x.shape[0], x.is_quantized) ri, zi, ni = self.igates(x).chunk(3, -1) if x.dim() > 2: for k in range(x.size(0)): hx = hidden rh, zh, nh = self.hgates(hx).chunk(3, -1) rgate = self.act1(self.add1.add(ri[k, ...], rh)) zgate = self.act1(self.add2.add(zi[k, ...], zh)) ngate = self.act2(self.add3.add(ni[k, ...], self.mul1.mul(rgate, nh))) hidden = self.add4.add( self.mul2.mul(self.sub1.add_scalar(self.mul4.mul_scalar(zgate, -1), 1), ngate), self.mul3.mul(zgate, hx), ) result.append(hidden) result_tensor = torch.stack(result, 0) return result_tensor, hidden else: log.warning('Make sure you are not passing unbatched input to GRU, which may yield errors.') hx = hidden rh, zh, nh = self.hgates(hx).chunk(3, -1) rgate = self.act1(self.add1.add(ri, rh)) zgate = self.act1(self.add2.add(zi, zh)) ngate = self.act2(self.add3.add(ni, self.mul1.mul(rgate, nh))) hidden = self.add4.add( self.mul2.mul(self.sub1.add_scalar(self.mul4.mul_scalar(zgate, -1), 1), ngate), self.mul3.mul(zgate, hx), ) result.append(hidden) result_tensor = torch.stack(result, 0) return result_tensor def initialize_hidden(self, batch_size: int, is_quantized: bool = False) -> Tensor: h = torch.zeros(batch_size, self.hidden_size) if is_quantized: (h_scale, h_zp) = self.initial_hidden_state_qparams h = torch.quantize_per_tensor(h, scale=h_scale, zero_point=h_zp, dtype=self.hidden_state_dtype) return h def _get_name(self): return 'QuantizableGRUCell' @classmethod def from_params(cls, wi, wh, bi=None, bh=None): """Uses the weights and biases to create a new GRU cell. Args: wi, wh: Weights for the input and hidden layers bi, bh: Biases for the input and hidden layers """ assert (bi is None) == (bh is None) # Either both None or both have values input_size = wi.shape[1] hidden_size = wh.shape[1] cell = cls(input_dim=input_size, hidden_dim=hidden_size, bias=(bi is not None)) cell.igates.weight = torch.nn.Parameter(wi) if bi is not None: cell.igates.bias = torch.nn.Parameter(bi) cell.hgates.weight = torch.nn.Parameter(wh) if bh is not None: cell.hgates.bias = torch.nn.Parameter(bh) return cell @classmethod def from_float(cls, other): assert type(other) is cls._FLOAT_MODULE assert hasattr(other, 'qconfig'), "The float module must have 'qconfig'" observed = cls.from_params(other.weight_ih, other.weight_hh, other.bias_ih, other.bias_hh) observed.qconfig = other.qconfig observed.igates.qconfig = other.qconfig observed.hgates.qconfig = other.qconfig return observed class _GRUSingleLayer(torch.nn.Module): r"""A single one-directional GRU layer. The difference between a layer and a cell is that the layer can process a sequence, while the cell only expects an instantaneous value. """ def __init__(self, input_dim: int, hidden_dim: int, bias: bool = True, device=None, dtype=None) -> None: factory_kwargs = {'device': device, 'dtype': dtype} super().__init__() self.cell = GRUCell(input_dim, hidden_dim, bias=bias, **factory_kwargs) def forward(self, x: Tensor, hidden: Optional[Tensor] = None): result_tensor, hidden = self.cell(x, hidden) return result_tensor, hidden @classmethod def from_params(cls, *args, **kwargs): cell = GRUCell.from_params(*args, **kwargs) layer = cls(cell.input_size, cell.hidden_size, cell.bias) layer.cell = cell return layer class _GRULayer(torch.nn.Module): r"""A single bi-directional GRU layer.""" def __init__( self, input_dim: int, hidden_dim: int, bias: bool = True, batch_first: bool = False, bidirectional: bool = False, device=None, dtype=None, ) -> None: factory_kwargs = {'device': device, 'dtype': dtype} super().__init__() self.batch_first = batch_first self.bidirectional = bidirectional self.layer_fw = _GRUSingleLayer(input_dim, hidden_dim, bias=bias, **factory_kwargs) if self.bidirectional: self.layer_bw = _GRUSingleLayer(input_dim, hidden_dim, bias=bias, **factory_kwargs) def forward(self, x: Tensor, hidden: Optional[Tensor] = None): if self.batch_first: x = x.transpose(0, 1) hx_fw = hidden hidden_bw: Optional[Tensor] = None if self.bidirectional: if hx_fw is None: hx_bw = None else: hx_bw = hx_fw[1] hx_fw = hx_fw[0] if hx_bw is not None: hidden_bw = hx_bw if hx_fw is None: hidden_fw = None else: hidden_fw = torch.jit._unwrap_optional(hx_fw) result_fw, hidden_fw = self.layer_fw(x, hidden_fw) if hasattr(self, "layer_bw") and self.bidirectional: x_reversed = x.flip(0) result_bw, hidden_bw = self.layer_bw(x_reversed, hidden_bw) result_bw = result_bw.flip(0) result = torch.cat([result_fw, result_bw], result_fw.dim() - 1) if hidden_fw is None and hidden_bw is None: h = None elif hidden_fw is None: h = torch.jit._unwrap_optional(hidden_fw) elif hidden_bw is None: h = torch.jit._unwrap_optional(hidden_bw) else: h = torch.stack([hidden_fw[0], hidden_bw[0]], 0) # type: ignore[list-item] else: result = result_fw h = torch.jit._unwrap_optional(hidden_fw) # type: ignore[assignment] if self.batch_first: result.transpose_(0, 1) return result, h @classmethod def from_float(cls, other, layer_idx=0, qconfig=None, **kwargs): r""" There is no FP equivalent of this class. This function is here just to mimic the behavior of the `prepare` within the `torch.ao.quantization` flow. """ assert hasattr(other, 'qconfig') or (qconfig is not None) input_size = kwargs.get('input_size', other.input_size) hidden_size = kwargs.get('hidden_size', other.hidden_size) bias = kwargs.get('bias', other.bias) batch_first = kwargs.get('batch_first', other.batch_first) bidirectional = kwargs.get('bidirectional', other.bidirectional) layer = cls(input_size, hidden_size, bias, batch_first, bidirectional) layer.qconfig = getattr(other, 'qconfig', qconfig) wi = getattr(other, f'weight_ih_l{layer_idx}') wh = getattr(other, f'weight_hh_l{layer_idx}') bi = getattr(other, f'bias_ih_l{layer_idx}', None) bh = getattr(other, f'bias_hh_l{layer_idx}', None) layer.layer_fw = _GRUSingleLayer.from_params(wi, wh, bi, bh) if other.bidirectional: wi = getattr(other, f'weight_ih_l{layer_idx}_reverse') wh = getattr(other, f'weight_hh_l{layer_idx}_reverse') bi = getattr(other, f'bias_ih_l{layer_idx}_reverse', None) bh = getattr(other, f'bias_hh_l{layer_idx}_reverse', None) layer.layer_bw = _GRUSingleLayer.from_params(wi, wh, bi, bh) return layer class GRU(torch.nn.Module): r"""A quantizable gated recurrent unit (GRU). For the description and the argument types, please, refer to :class:`~torch.nn.GRU` Attributes: layers : instances of the `_GRULayer` .. note:: To access the weights and biases, you need to access them per layer. AssertionError: There is no reverse path in the non-bidirectional layer """ _FLOAT_MODULE = torch.nn.GRU def __init__( self, input_size: int, hidden_size: int, num_layers: int = 1, bias: bool = True, batch_first: bool = False, dropout: float = 0.0, bidirectional: bool = False, device=None, dtype=None, ) -> None: factory_kwargs = {'device': device, 'dtype': dtype} super().__init__() self.input_size = input_size self.hidden_size = hidden_size self.num_layers = num_layers self.bias = bias self.batch_first = batch_first self.dropout = float(dropout) self.bidirectional = bidirectional self.training = False # We don't want to train using this module if not isinstance(dropout, numbers.Number) or not 0 <= dropout <= 1 or isinstance(dropout, bool): raise ValueError( "dropout should be a number in range [0, 1] representing the probability of an element being zeroed" ) if dropout > 0: warnings.warn( "dropout option for quantizable GRU is ignored. " "If you are training, please, use nn.GRU version " "followed by `prepare` step." ) if num_layers == 1: warnings.warn( "dropout option adds dropout after all but last " "recurrent layer, so non-zero dropout expects " "num_layers greater than 1, but got dropout={} " "and num_layers={}".format(dropout, num_layers) ) layers = [ _GRULayer( self.input_size, self.hidden_size, self.bias, batch_first=False, bidirectional=self.bidirectional, **factory_kwargs, ) ] for layer in range(1, num_layers): layers.append( _GRULayer( self.hidden_size, self.hidden_size, self.bias, batch_first=False, bidirectional=self.bidirectional, **factory_kwargs, ) ) self.layers = torch.nn.ModuleList(layers) def forward(self, x: Tensor, hidden: Optional[Tensor] = None): if self.batch_first: x = x.transpose(0, 1) max_batch_size = x.size(1) num_directions = 2 if self.bidirectional else 1 if hidden is None: zeros = torch.zeros( num_directions, max_batch_size, self.hidden_size, dtype=torch.float, device=x.device ) zeros.squeeze_(0) if x.is_quantized: zeros = torch.quantize_per_tensor(zeros, scale=1.0, zero_point=0, dtype=x.dtype) hx = [zeros for _ in range(self.num_layers)] else: hidden_non_opt = torch.jit._unwrap_optional(hidden) if isinstance(hidden_non_opt, Tensor): hx = hidden_non_opt.reshape( self.num_layers, num_directions, max_batch_size, self.hidden_size ).unbind(0) hx = [(hx[idx].squeeze_(0)) for idx in range(self.num_layers)] elif isinstance(hidden_non_opt[0], Tensor): hx = ( hidden_non_opt[0] .reshape(self.num_layers, num_directions, max_batch_size, self.hidden_size) .unbind(0) ) hx = [(hx[idx].squeeze_(0)) for idx in range(self.num_layers)] else: hx = hidden_non_opt hx_list = [] for idx, layer in enumerate(self.layers): x, h = layer(x, hx[idx]) hx_list.append(torch.jit._unwrap_optional(h)) hx_tensor = torch.stack(hx_list) # We are creating another dimension for bidirectional case # need to collapse it hx_tensor = hx_tensor.reshape(-1, hx_tensor.shape[-2], hx_tensor.shape[-1]) if self.batch_first: x = x.transpose(0, 1) return x, hx_tensor def _get_name(self): return 'QuantizableGRU' @classmethod def from_float(cls, other, qconfig=None): assert isinstance(other, cls._FLOAT_MODULE) assert hasattr(other, 'qconfig') or qconfig observed = cls( other.input_size, other.hidden_size, other.num_layers, other.bias, other.batch_first, other.dropout, other.bidirectional, ) observed.qconfig = getattr(other, 'qconfig', qconfig) for idx in range(other.num_layers): observed.layers[idx] = _GRULayer.from_float(other, idx, qconfig, batch_first=False) observed.eval() observed = torch.ao.quantization.prepare(observed, inplace=True) return observed @classmethod def from_observed(cls, other): # The whole flow is float -> observed -> quantized # This class does float -> observed only raise NotImplementedError( "It looks like you are trying to convert a " "non-quantizable GRU module. Please, see " "the examples on quantizable GRUs." )