Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1 @@
|
||||
from .modules import * # noqa: F403
|
||||
@@ -0,0 +1,9 @@
|
||||
from .activation import MultiheadAttention
|
||||
from .rnn import LSTM, LSTMCell
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LSTM",
|
||||
"LSTMCell",
|
||||
"MultiheadAttention",
|
||||
]
|
||||
+623
@@ -0,0 +1,623 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
import torch.jit # this is needed to avoid a circular import
|
||||
import torch.nn.functional as F
|
||||
from torch import nn, Tensor
|
||||
|
||||
|
||||
__all__ = ["MultiheadAttention"]
|
||||
|
||||
|
||||
class MultiheadAttention(nn.MultiheadAttention):
|
||||
_FLOAT_MODULE = nn.MultiheadAttention
|
||||
|
||||
r"""Quantizable implementation of the MultiheadAttention.
|
||||
|
||||
Note::
|
||||
Please, refer to :class:`~torch.nn.MultiheadAttention` for more
|
||||
information
|
||||
|
||||
Allows the model to jointly attend to information from different
|
||||
representation subspaces.
|
||||
See reference: Attention Is All You Need
|
||||
|
||||
The original MHA module is not quantizable.
|
||||
This reimplements it by explicitly instantiating the linear layers.
|
||||
|
||||
.. math::
|
||||
\text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
|
||||
\text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
|
||||
|
||||
Args:
|
||||
embed_dim: total dimension of the model.
|
||||
num_heads: parallel attention heads.
|
||||
dropout: a Dropout layer on attn_output_weights. Default: 0.0.
|
||||
bias: add bias as module parameter. Default: True.
|
||||
add_bias_kv: add bias to the key and value sequences at dim=0.
|
||||
add_zero_attn: add a new batch of zeros to the key and
|
||||
value sequences at dim=1.
|
||||
kdim: total number of features in key. Default: None.
|
||||
vdim: total number of features in value. Default: None.
|
||||
batch_first: If ``True``, then the input and output tensors are provided
|
||||
as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
|
||||
|
||||
Note that if :attr:`kdim` and :attr:`vdim` are None, they will be set
|
||||
to :attr:`embed_dim` such that query, key, and value have the same
|
||||
number of features.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> multihead_attn = nnqa.MultiheadAttention(embed_dim, num_heads)
|
||||
>>> attn_output, attn_output_weights = multihead_attn(query, key, value)
|
||||
|
||||
Note::
|
||||
Please, follow the quantization flow to convert the quantizable MHA.
|
||||
"""
|
||||
__constants__ = ["batch_first"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embed_dim: int,
|
||||
num_heads: int,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = True,
|
||||
add_bias_kv: bool = False,
|
||||
add_zero_attn: bool = False,
|
||||
kdim: int | None = None,
|
||||
vdim: int | None = None,
|
||||
batch_first: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__(
|
||||
embed_dim,
|
||||
num_heads,
|
||||
dropout,
|
||||
bias,
|
||||
add_bias_kv,
|
||||
add_zero_attn,
|
||||
kdim,
|
||||
vdim,
|
||||
batch_first,
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.linear_Q = nn.Linear(
|
||||
self.embed_dim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.linear_K = nn.Linear(
|
||||
self.kdim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.linear_V = nn.Linear(
|
||||
self.vdim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# for the type: ignore, see https://github.com/pytorch/pytorch/issues/58969
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.out_proj = nn.Linear(
|
||||
self.embed_dim, self.embed_dim, bias=bias, **factory_kwargs
|
||||
) # type: ignore[assignment]
|
||||
|
||||
# Functionals
|
||||
self.q_scaling_product = torch.ao.nn.quantized.FloatFunctional()
|
||||
# note: importing torch.ao.nn.quantized at top creates a circular import
|
||||
|
||||
# Quant/Dequant
|
||||
self.quant_attn_output = torch.ao.quantization.QuantStub()
|
||||
self.quant_attn_output_weights = torch.ao.quantization.QuantStub()
|
||||
self.dequant_q = torch.ao.quantization.DeQuantStub()
|
||||
self.dequant_k = torch.ao.quantization.DeQuantStub()
|
||||
self.dequant_v = torch.ao.quantization.DeQuantStub()
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableMultiheadAttention"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other):
|
||||
if type(other) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"Expected type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig"):
|
||||
raise AssertionError("The float module must have 'qconfig'")
|
||||
# Setting the dropout to 0.0!
|
||||
observed = cls(
|
||||
other.embed_dim,
|
||||
other.num_heads,
|
||||
other.dropout,
|
||||
(other.in_proj_bias is not None),
|
||||
(other.bias_k is not None),
|
||||
other.add_zero_attn,
|
||||
other.kdim,
|
||||
other.vdim,
|
||||
other.batch_first,
|
||||
)
|
||||
observed.bias_k = other.bias_k
|
||||
observed.bias_v = other.bias_v
|
||||
observed.qconfig = other.qconfig
|
||||
|
||||
# Set the linear weights
|
||||
# for the type: ignores, see https://github.com/pytorch/pytorch/issues/58969
|
||||
observed.out_proj.weight = other.out_proj.weight
|
||||
observed.out_proj.bias = other.out_proj.bias
|
||||
if other._qkv_same_embed_dim:
|
||||
# Use separate params
|
||||
bias = other.in_proj_bias
|
||||
_start = 0
|
||||
_end = _start + other.embed_dim
|
||||
weight = other.in_proj_weight[_start:_end, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:_end], bias.requires_grad)
|
||||
observed.linear_Q.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_Q.bias = bias
|
||||
|
||||
bias = other.in_proj_bias
|
||||
_start = _end
|
||||
_end = _start + other.embed_dim
|
||||
weight = other.in_proj_weight[_start:_end, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:_end], bias.requires_grad)
|
||||
observed.linear_K.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_K.bias = bias
|
||||
|
||||
bias = other.in_proj_bias
|
||||
_start = _end
|
||||
weight = other.in_proj_weight[_start:, :]
|
||||
if bias is not None:
|
||||
bias = torch.nn.Parameter(bias[_start:], bias.requires_grad)
|
||||
observed.linear_V.weight = torch.nn.Parameter(weight, weight.requires_grad)
|
||||
observed.linear_V.bias = bias
|
||||
else:
|
||||
observed.linear_Q.weight = nn.Parameter(other.q_proj_weight)
|
||||
observed.linear_K.weight = nn.Parameter(other.k_proj_weight)
|
||||
observed.linear_V.weight = nn.Parameter(other.v_proj_weight)
|
||||
if other.in_proj_bias is None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_Q.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_K.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
observed.linear_V.bias = None
|
||||
else:
|
||||
observed.linear_Q.bias = nn.Parameter(
|
||||
other.in_proj_bias[0 : other.embed_dim]
|
||||
)
|
||||
observed.linear_K.bias = nn.Parameter(
|
||||
other.in_proj_bias[other.embed_dim : (other.embed_dim * 2)]
|
||||
)
|
||||
observed.linear_V.bias = nn.Parameter(
|
||||
other.in_proj_bias[(other.embed_dim * 2) :]
|
||||
)
|
||||
observed.eval()
|
||||
# Explicit prepare
|
||||
observed = torch.ao.quantization.prepare(observed, inplace=True)
|
||||
return observed
|
||||
|
||||
@torch.jit.unused
|
||||
def dequantize(self):
|
||||
r"""Utility to convert the quantized MHA back to float.
|
||||
|
||||
The motivation for this is that it is not trivial to convert the weights
|
||||
from the format that is used in the quantized version back to the
|
||||
float.
|
||||
"""
|
||||
fp = self._FLOAT_MODULE(
|
||||
self.embed_dim,
|
||||
self.num_heads,
|
||||
self.dropout,
|
||||
(self.linear_Q._weight_bias()[1] is not None), # type: ignore[operator]
|
||||
(self.bias_k is not None),
|
||||
self.add_zero_attn,
|
||||
self.kdim,
|
||||
self.vdim,
|
||||
self.batch_first,
|
||||
)
|
||||
if fp._qkv_same_embed_dim != self._qkv_same_embed_dim:
|
||||
raise AssertionError(
|
||||
f"_qkv_same_embed_dim mismatch: {fp._qkv_same_embed_dim} != {self._qkv_same_embed_dim}"
|
||||
)
|
||||
if self.bias_k is not None:
|
||||
fp.bias_k = nn.Parameter(self.bias_k.dequantize())
|
||||
if self.bias_v is not None:
|
||||
fp.bias_v = nn.Parameter(self.bias_v.dequantize())
|
||||
|
||||
# Set the linear weights
|
||||
# Note: Because the linear layers are quantized, mypy does not know how
|
||||
# to deal with them -- might need to ignore the typing checks.
|
||||
# for the type: ignore[has-type], see https://github.com/pytorch/pytorch/issues/58969
|
||||
w, b = self.out_proj._weight_bias() # type: ignore[operator, has-type]
|
||||
fp.out_proj.weight = nn.Parameter(w.dequantize())
|
||||
if b is not None:
|
||||
fp.out_proj.bias = nn.Parameter(b)
|
||||
|
||||
wQ, bQ = self.linear_Q._weight_bias() # type: ignore[operator]
|
||||
wQ = wQ.dequantize()
|
||||
wK, bK = self.linear_K._weight_bias() # type: ignore[operator]
|
||||
wK = wK.dequantize()
|
||||
wV, bV = self.linear_V._weight_bias() # type: ignore[operator]
|
||||
wV = wV.dequantize()
|
||||
if fp._qkv_same_embed_dim:
|
||||
# Use separate params
|
||||
_start = 0
|
||||
_end = _start + fp.embed_dim
|
||||
fp.in_proj_weight[_start:_end, :] = wQ
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bQ == 0):
|
||||
raise AssertionError("Expected all bQ elements to be 0")
|
||||
fp.in_proj_bias[_start:_end] = bQ
|
||||
|
||||
_start = _end
|
||||
_end = _start + fp.embed_dim
|
||||
fp.in_proj_weight[_start:_end, :] = wK
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bK == 0):
|
||||
raise AssertionError("Expected all bK elements to be 0")
|
||||
fp.in_proj_bias[_start:_end] = bK
|
||||
|
||||
_start = _end
|
||||
fp.in_proj_weight[_start:, :] = wV
|
||||
if fp.in_proj_bias is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
if not all(bV == 0):
|
||||
raise AssertionError("Expected all bV elements to be 0")
|
||||
fp.in_proj_bias[_start:] = bV
|
||||
else:
|
||||
fp.q_proj_weight = nn.Parameter(wQ)
|
||||
fp.k_proj_weight = nn.Parameter(wK)
|
||||
fp.v_proj_weight = nn.Parameter(wV)
|
||||
if fp.in_proj_bias is None:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_Q.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_K.bias = None
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.linear_V.bias = None
|
||||
else:
|
||||
fp.in_proj_bias[0 : fp.embed_dim] = bQ
|
||||
fp.in_proj_bias[fp.embed_dim : (fp.embed_dim * 2)] = bK
|
||||
fp.in_proj_bias[(fp.embed_dim * 2) :] = bV
|
||||
|
||||
return fp
|
||||
|
||||
@classmethod
|
||||
def from_observed(cls, other):
|
||||
# The whole flow is float -> observed -> quantized
|
||||
# This class does float -> observed only
|
||||
# See nn.quantized.MultiheadAttention
|
||||
raise NotImplementedError(
|
||||
"It looks like you are trying to prepare an "
|
||||
"MHA module. Please, see "
|
||||
"the examples on quantizable MHAs."
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
key_padding_mask: Tensor | None = None,
|
||||
need_weights: bool = True,
|
||||
attn_mask: Tensor | None = None,
|
||||
average_attn_weights: bool = True,
|
||||
is_causal: bool = False,
|
||||
) -> tuple[Tensor, Tensor | None]:
|
||||
r"""
|
||||
Note::
|
||||
Please, refer to :func:`~torch.nn.MultiheadAttention.forward` for more
|
||||
information
|
||||
|
||||
Args:
|
||||
query, key, value: map a query and a set of key-value pairs to an output.
|
||||
See "Attention Is All You Need" for more details.
|
||||
key_padding_mask: if provided, specified padding elements in the key will
|
||||
be ignored by the attention. When given a binary mask and a value is True,
|
||||
the corresponding value on the attention layer will be ignored.
|
||||
need_weights: output attn_output_weights.
|
||||
attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
|
||||
the batches while a 3D mask allows to specify a different mask for the entries of each batch.
|
||||
|
||||
Shape:
|
||||
- Inputs:
|
||||
- query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
|
||||
- key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
|
||||
- value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
|
||||
the embedding dimension. :math:`(N, S, E)` if ``batch_first`` is ``True``.
|
||||
- key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
|
||||
If a BoolTensor is provided, the positions with the
|
||||
value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
|
||||
- attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
|
||||
3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
|
||||
S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
|
||||
positions. If a BoolTensor is provided, positions with ``True``
|
||||
is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
|
||||
is provided, it will be added to the attention weight.
|
||||
- is_causal: If specified, applies a causal mask as attention mask. Mutually exclusive with providing attn_mask.
|
||||
Default: ``False``.
|
||||
- average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
|
||||
heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
|
||||
effect when ``need_weights=True.``. Default: True (i.e. average weights across heads)
|
||||
|
||||
- Outputs:
|
||||
- attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
|
||||
E is the embedding dimension. :math:`(N, L, E)` if ``batch_first`` is ``True``.
|
||||
- attn_output_weights: If ``average_attn_weights=True``, returns attention weights averaged
|
||||
across heads of shape :math:`(N, L, S)`, where N is the batch size, L is the target sequence length,
|
||||
S is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
|
||||
head of shape :math:`(N, num_heads, L, S)`.
|
||||
"""
|
||||
return self._forward_impl(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
key_padding_mask,
|
||||
need_weights,
|
||||
attn_mask,
|
||||
average_attn_weights,
|
||||
is_causal,
|
||||
)
|
||||
|
||||
def _forward_impl(
|
||||
self,
|
||||
query: Tensor,
|
||||
key: Tensor,
|
||||
value: Tensor,
|
||||
key_padding_mask: Tensor | None = None,
|
||||
need_weights: bool = True,
|
||||
attn_mask: Tensor | None = None,
|
||||
average_attn_weights: bool = True,
|
||||
is_causal: bool = False,
|
||||
) -> tuple[Tensor, Tensor | None]:
|
||||
# This version will not deal with the static key/value pairs.
|
||||
# Keeping it here for future changes.
|
||||
#
|
||||
# TODO: This method has some duplicate lines with the
|
||||
# `torch.nn.functional.multi_head_attention`. Will need to refactor.
|
||||
static_k = None
|
||||
static_v = None
|
||||
|
||||
if attn_mask is not None and is_causal:
|
||||
raise AssertionError("Only allow causal mask or attn_mask")
|
||||
|
||||
if is_causal:
|
||||
raise AssertionError("causal mask not supported by AO MHA module")
|
||||
|
||||
if self.batch_first:
|
||||
query, key, value = (x.transpose(0, 1) for x in (query, key, value))
|
||||
|
||||
tgt_len, bsz, embed_dim_to_check = query.size()
|
||||
if self.embed_dim != embed_dim_to_check:
|
||||
raise AssertionError(
|
||||
f"embed_dim mismatch: {self.embed_dim} != {embed_dim_to_check}"
|
||||
)
|
||||
# allow MHA to have different sizes for the feature dimension
|
||||
if key.size(0) != value.size(0) or key.size(1) != value.size(1):
|
||||
raise AssertionError(
|
||||
f"key and value size mismatch: key.size()={key.size()}, value.size()={value.size()}"
|
||||
)
|
||||
|
||||
head_dim = self.embed_dim // self.num_heads
|
||||
if head_dim * self.num_heads != self.embed_dim:
|
||||
raise AssertionError("embed_dim must be divisible by num_heads")
|
||||
scaling = float(head_dim) ** -0.5
|
||||
|
||||
q = self.linear_Q(query)
|
||||
k = self.linear_K(key)
|
||||
v = self.linear_V(value)
|
||||
|
||||
q = self.q_scaling_product.mul_scalar(q, scaling)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.uint8:
|
||||
warnings.warn(
|
||||
"Byte tensor for `attn_mask` in `nn.MultiheadAttention` is deprecated. "
|
||||
"Use bool tensor instead.",
|
||||
stacklevel=3,
|
||||
)
|
||||
attn_mask = attn_mask.to(torch.bool)
|
||||
if not attn_mask.is_floating_point() and attn_mask.dtype != torch.bool:
|
||||
raise AssertionError(
|
||||
f"Only float and bool types are supported for attn_mask, not {attn_mask.dtype}"
|
||||
)
|
||||
|
||||
if attn_mask.dim() == 2:
|
||||
attn_mask = attn_mask.unsqueeze(0)
|
||||
if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
|
||||
raise RuntimeError("The size of the 2D attn_mask is not correct.")
|
||||
elif attn_mask.dim() == 3:
|
||||
if list(attn_mask.size()) != [
|
||||
bsz * self.num_heads,
|
||||
query.size(0),
|
||||
key.size(0),
|
||||
]:
|
||||
raise RuntimeError("The size of the 3D attn_mask is not correct.")
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"attn_mask's dimension {attn_mask.dim()} is not supported"
|
||||
)
|
||||
# attn_mask's dim is 3 now.
|
||||
|
||||
# convert ByteTensor key_padding_mask to bool
|
||||
if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
|
||||
warnings.warn(
|
||||
"Byte tensor for `key_padding_mask` in `nn.MultiheadAttention` is deprecated. "
|
||||
"Use bool tensor instead.",
|
||||
stacklevel=3,
|
||||
)
|
||||
key_padding_mask = key_padding_mask.to(torch.bool)
|
||||
if self.bias_k is not None and self.bias_v is not None:
|
||||
if static_k is None and static_v is None:
|
||||
# Explicitly check that bias_k and bias_v are not None
|
||||
# in a way that TorchScript can understand.
|
||||
bias_k = self.bias_k
|
||||
if bias_k is None:
|
||||
raise AssertionError("bias_k must not be None")
|
||||
bias_v = self.bias_v
|
||||
if bias_v is None:
|
||||
raise AssertionError("bias_v must not be None")
|
||||
|
||||
k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
|
||||
v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
|
||||
if attn_mask is not None:
|
||||
attn_mask = F.pad(attn_mask, (0, 1))
|
||||
if key_padding_mask is not None:
|
||||
key_padding_mask = F.pad(key_padding_mask, (0, 1))
|
||||
else:
|
||||
if static_k is not None:
|
||||
raise AssertionError("bias cannot be added to static key.")
|
||||
if static_v is not None:
|
||||
raise AssertionError("bias cannot be added to static value.")
|
||||
else:
|
||||
if self.bias_k is not None:
|
||||
raise AssertionError(
|
||||
"self.bias_k must be None when self.bias_v is None"
|
||||
)
|
||||
if self.bias_v is not None:
|
||||
raise AssertionError(
|
||||
"self.bias_v must be None when self.bias_k is None"
|
||||
)
|
||||
|
||||
q = q.contiguous().view(tgt_len, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
if k is not None:
|
||||
k = k.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
if v is not None:
|
||||
v = v.contiguous().view(-1, bsz * self.num_heads, head_dim).transpose(0, 1)
|
||||
|
||||
if static_k is not None:
|
||||
if static_k.size(0) != bsz * self.num_heads:
|
||||
raise AssertionError(
|
||||
f"static_k.size(0) must be {bsz * self.num_heads}, got {static_k.size(0)}"
|
||||
)
|
||||
if static_k.size(2) != head_dim:
|
||||
raise AssertionError(
|
||||
f"static_k.size(2) must be {head_dim}, got {static_k.size(2)}"
|
||||
)
|
||||
k = static_k
|
||||
|
||||
if static_v is not None:
|
||||
if static_v.size(0) != bsz * self.num_heads:
|
||||
raise AssertionError(
|
||||
f"static_v.size(0) must be {bsz * self.num_heads}, got {static_v.size(0)}"
|
||||
)
|
||||
if static_v.size(2) != head_dim:
|
||||
raise AssertionError(
|
||||
f"static_v.size(2) must be {head_dim}, got {static_v.size(2)}"
|
||||
)
|
||||
v = static_v
|
||||
|
||||
src_len = k.size(1)
|
||||
|
||||
if key_padding_mask is not None:
|
||||
if key_padding_mask.size(0) != bsz:
|
||||
raise AssertionError(
|
||||
f"key_padding_mask.size(0) must be {bsz}, got {key_padding_mask.size(0)}"
|
||||
)
|
||||
if key_padding_mask.size(1) != src_len:
|
||||
raise AssertionError(
|
||||
f"key_padding_mask.size(1) must be {src_len}, got {key_padding_mask.size(1)}"
|
||||
)
|
||||
|
||||
if self.add_zero_attn:
|
||||
src_len += 1
|
||||
|
||||
k_zeros = torch.zeros((k.size(0), 1) + k.size()[2:])
|
||||
|
||||
if k.is_quantized:
|
||||
k_zeros = torch.quantize_per_tensor(
|
||||
k_zeros,
|
||||
k.q_scale(),
|
||||
k.q_zero_point(),
|
||||
k.dtype,
|
||||
)
|
||||
|
||||
k = torch.cat([k, k_zeros], dim=1)
|
||||
|
||||
v_zeros = torch.zeros((v.size(0), 1) + k.size()[2:])
|
||||
|
||||
if v.is_quantized:
|
||||
v_zeros = torch.quantize_per_tensor(
|
||||
v_zeros,
|
||||
v.q_scale(),
|
||||
v.q_zero_point(),
|
||||
v.dtype,
|
||||
)
|
||||
|
||||
v = torch.cat([v, v_zeros], dim=1)
|
||||
|
||||
if attn_mask is not None:
|
||||
attn_mask = F.pad(attn_mask, (0, 1))
|
||||
if key_padding_mask is not None:
|
||||
key_padding_mask = F.pad(key_padding_mask, (0, 1))
|
||||
|
||||
# Leaving the quantized zone here
|
||||
q = self.dequant_q(q)
|
||||
k = self.dequant_k(k)
|
||||
v = self.dequant_v(v)
|
||||
attn_output_weights = torch.bmm(q, k.transpose(1, 2))
|
||||
expected_size = [bsz * self.num_heads, tgt_len, src_len]
|
||||
if list(attn_output_weights.size()) != expected_size:
|
||||
raise AssertionError(
|
||||
f"attn_output_weights size mismatch: expected {expected_size}, "
|
||||
f"got {list(attn_output_weights.size())}"
|
||||
)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_output_weights.masked_fill_(attn_mask, float("-inf"))
|
||||
else:
|
||||
attn_output_weights += attn_mask
|
||||
|
||||
if key_padding_mask is not None:
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz, self.num_heads, tgt_len, src_len
|
||||
)
|
||||
attn_output_weights = attn_output_weights.masked_fill(
|
||||
key_padding_mask.unsqueeze(1).unsqueeze(2),
|
||||
float("-inf"),
|
||||
)
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz * self.num_heads, tgt_len, src_len
|
||||
)
|
||||
|
||||
attn_output_weights = F.softmax(attn_output_weights, dim=-1)
|
||||
attn_output_weights = F.dropout(
|
||||
attn_output_weights, p=self.dropout, training=self.training
|
||||
)
|
||||
|
||||
attn_output = torch.bmm(attn_output_weights, v)
|
||||
expected_output_size = [bsz * self.num_heads, tgt_len, head_dim]
|
||||
if list(attn_output.size()) != expected_output_size:
|
||||
raise AssertionError(
|
||||
f"attn_output size mismatch: expected {expected_output_size}, "
|
||||
f"got {list(attn_output.size())}"
|
||||
)
|
||||
if self.batch_first:
|
||||
attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
|
||||
else:
|
||||
attn_output = (
|
||||
attn_output.transpose(0, 1)
|
||||
.contiguous()
|
||||
.view(tgt_len, bsz, self.embed_dim)
|
||||
)
|
||||
|
||||
# Reentering the quantized zone
|
||||
attn_output = self.quant_attn_output(attn_output)
|
||||
# for the type: ignore[has-type], see https://github.com/pytorch/pytorch/issues/58969
|
||||
attn_output = self.out_proj(attn_output) # type: ignore[has-type]
|
||||
attn_output_weights = self.quant_attn_output_weights(attn_output_weights)
|
||||
|
||||
if need_weights:
|
||||
# average attention weights over heads
|
||||
attn_output_weights = attn_output_weights.view(
|
||||
bsz, self.num_heads, tgt_len, src_len
|
||||
)
|
||||
if average_attn_weights:
|
||||
attn_output_weights = attn_output_weights.mean(dim=1)
|
||||
return attn_output, attn_output_weights
|
||||
else:
|
||||
return attn_output, None
|
||||
@@ -0,0 +1,618 @@
|
||||
"""
|
||||
We will recreate all the RNN modules as we require the modules to be decomposed
|
||||
into its building blocks to be able to observe.
|
||||
"""
|
||||
|
||||
# mypy: allow-untyped-defs
|
||||
|
||||
import numbers
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
__all__ = ["LSTMCell", "LSTM"]
|
||||
|
||||
|
||||
class LSTMCell(torch.nn.Module):
|
||||
r"""A quantizable long short-term memory (LSTM) cell.
|
||||
|
||||
For the description and the argument types, please, refer to :class:`~torch.nn.LSTMCell`
|
||||
|
||||
`split_gates`: specify True to compute the input/forget/cell/output gates separately
|
||||
to avoid an intermediate tensor which is subsequently chunk'd. This optimization can
|
||||
be beneficial for on-device inference latency. This flag is cascaded down from the
|
||||
parent classes.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> rnn = nnqa.LSTMCell(10, 20)
|
||||
>>> input = torch.randn(6, 10)
|
||||
>>> hx = torch.randn(3, 20)
|
||||
>>> cx = torch.randn(3, 20)
|
||||
>>> output = []
|
||||
>>> for i in range(6):
|
||||
... hx, cx = rnn(input[i], (hx, cx))
|
||||
... output.append(hx)
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.nn.LSTMCell
|
||||
__constants__ = ["split_gates"] # for jit.script
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
bias: bool = True,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.input_size = input_dim
|
||||
self.hidden_size = hidden_dim
|
||||
self.bias = bias
|
||||
self.split_gates = split_gates
|
||||
|
||||
if not split_gates:
|
||||
self.igates: torch.nn.Module = torch.nn.Linear(
|
||||
input_dim, 4 * hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.hgates: torch.nn.Module = torch.nn.Linear(
|
||||
hidden_dim, 4 * hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
self.gates: torch.nn.Module = torch.ao.nn.quantized.FloatFunctional()
|
||||
else:
|
||||
# keep separate Linear layers for each gate
|
||||
self.igates = torch.nn.ModuleDict()
|
||||
self.hgates = torch.nn.ModuleDict()
|
||||
self.gates = torch.nn.ModuleDict()
|
||||
for g in ["input", "forget", "cell", "output"]:
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.igates[g] = torch.nn.Linear(
|
||||
input_dim, hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.hgates[g] = torch.nn.Linear(
|
||||
hidden_dim, hidden_dim, bias=bias, **factory_kwargs
|
||||
)
|
||||
# pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
|
||||
self.gates[g] = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.input_gate = torch.nn.Sigmoid()
|
||||
self.forget_gate = torch.nn.Sigmoid()
|
||||
self.cell_gate = torch.nn.Tanh()
|
||||
self.output_gate = torch.nn.Sigmoid()
|
||||
|
||||
self.fgate_cx = torch.ao.nn.quantized.FloatFunctional()
|
||||
self.igate_cgate = torch.ao.nn.quantized.FloatFunctional()
|
||||
self.fgate_cx_igate_cgate = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.ogate_cy = torch.ao.nn.quantized.FloatFunctional()
|
||||
|
||||
self.initial_hidden_state_qparams: tuple[float, int] = (1.0, 0)
|
||||
self.initial_cell_state_qparams: tuple[float, int] = (1.0, 0)
|
||||
self.hidden_state_dtype: torch.dtype = torch.quint8
|
||||
self.cell_state_dtype: torch.dtype = torch.quint8
|
||||
|
||||
def forward(
|
||||
self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
if hidden is None or hidden[0] is None or hidden[1] is None:
|
||||
hidden = self.initialize_hidden(x.shape[0], x.is_quantized)
|
||||
hx, cx = hidden
|
||||
|
||||
if not self.split_gates:
|
||||
igates = self.igates(x)
|
||||
hgates = self.hgates(hx)
|
||||
gates = self.gates.add(igates, hgates) # type: ignore[operator]
|
||||
|
||||
input_gate, forget_gate, cell_gate, out_gate = gates.chunk(4, 1)
|
||||
|
||||
input_gate = self.input_gate(input_gate)
|
||||
forget_gate = self.forget_gate(forget_gate)
|
||||
cell_gate = self.cell_gate(cell_gate)
|
||||
out_gate = self.output_gate(out_gate)
|
||||
else:
|
||||
# apply each input + hidden projection and add together
|
||||
gate = {}
|
||||
for (key, gates), igates, hgates in zip(
|
||||
self.gates.items(), # type: ignore[operator]
|
||||
self.igates.values(), # type: ignore[operator]
|
||||
self.hgates.values(), # type: ignore[operator]
|
||||
):
|
||||
gate[key] = gates.add(igates(x), hgates(hx))
|
||||
|
||||
input_gate = self.input_gate(gate["input"])
|
||||
forget_gate = self.forget_gate(gate["forget"])
|
||||
cell_gate = self.cell_gate(gate["cell"])
|
||||
out_gate = self.output_gate(gate["output"])
|
||||
|
||||
fgate_cx = self.fgate_cx.mul(forget_gate, cx)
|
||||
igate_cgate = self.igate_cgate.mul(input_gate, cell_gate)
|
||||
fgate_cx_igate_cgate = self.fgate_cx_igate_cgate.add(fgate_cx, igate_cgate)
|
||||
cy = fgate_cx_igate_cgate
|
||||
|
||||
# TODO: make this tanh a member of the module so its qparams can be configured
|
||||
tanh_cy = torch.tanh(cy)
|
||||
hy = self.ogate_cy.mul(out_gate, tanh_cy)
|
||||
return hy, cy
|
||||
|
||||
def initialize_hidden(
|
||||
self, batch_size: int, is_quantized: bool = False
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
h, c = (
|
||||
torch.zeros((batch_size, self.hidden_size)),
|
||||
torch.zeros((batch_size, self.hidden_size)),
|
||||
)
|
||||
if is_quantized:
|
||||
(h_scale, h_zp) = self.initial_hidden_state_qparams
|
||||
(c_scale, c_zp) = self.initial_cell_state_qparams
|
||||
h = torch.quantize_per_tensor(
|
||||
h, scale=h_scale, zero_point=h_zp, dtype=self.hidden_state_dtype
|
||||
)
|
||||
c = torch.quantize_per_tensor(
|
||||
c, scale=c_scale, zero_point=c_zp, dtype=self.cell_state_dtype
|
||||
)
|
||||
return h, c
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableLSTMCell"
|
||||
|
||||
@classmethod
|
||||
def from_params(cls, wi, wh, bi=None, bh=None, split_gates=False):
|
||||
"""Uses the weights and biases to create a new LSTM cell.
|
||||
|
||||
Args:
|
||||
wi, wh: Weights for the input and hidden layers
|
||||
bi, bh: Biases for the input and hidden layers
|
||||
"""
|
||||
if (bi is None) != (bh is None):
|
||||
raise AssertionError("bi and bh must both be 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),
|
||||
split_gates=split_gates,
|
||||
)
|
||||
|
||||
if not split_gates:
|
||||
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)
|
||||
else:
|
||||
# split weight/bias
|
||||
for w, b, gates in zip([wi, wh], [bi, bh], [cell.igates, cell.hgates]):
|
||||
for w_chunk, gate in zip(w.chunk(4, dim=0), gates.values()): # type: ignore[operator]
|
||||
gate.weight = torch.nn.Parameter(w_chunk)
|
||||
|
||||
if b is not None:
|
||||
for b_chunk, gate in zip(b.chunk(4, dim=0), gates.values()): # type: ignore[operator]
|
||||
gate.bias = torch.nn.Parameter(b_chunk)
|
||||
|
||||
return cell
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other, use_precomputed_fake_quant=False, split_gates=False):
|
||||
if type(other) is not cls._FLOAT_MODULE:
|
||||
raise AssertionError(
|
||||
f"Expected module type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig"):
|
||||
raise AssertionError("The float module must have 'qconfig'")
|
||||
observed = cls.from_params(
|
||||
other.weight_ih,
|
||||
other.weight_hh,
|
||||
other.bias_ih,
|
||||
other.bias_hh,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
observed.qconfig = other.qconfig
|
||||
observed.igates.qconfig = other.qconfig
|
||||
observed.hgates.qconfig = other.qconfig
|
||||
if split_gates:
|
||||
# also apply qconfig directly to Linear modules
|
||||
for g in observed.igates.values():
|
||||
g.qconfig = other.qconfig
|
||||
for g in observed.hgates.values():
|
||||
g.qconfig = other.qconfig
|
||||
return observed
|
||||
|
||||
|
||||
class _LSTMSingleLayer(torch.nn.Module):
|
||||
r"""A single one-directional LSTM 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,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.cell = LSTMCell(
|
||||
input_dim, hidden_dim, bias=bias, split_gates=split_gates, **factory_kwargs
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
|
||||
result = []
|
||||
seq_len = x.shape[0]
|
||||
for i in range(seq_len):
|
||||
hidden = self.cell(x[i], hidden)
|
||||
result.append(hidden[0]) # type: ignore[index]
|
||||
result_tensor = torch.stack(result, 0)
|
||||
return result_tensor, hidden
|
||||
|
||||
@classmethod
|
||||
def from_params(cls, *args, **kwargs):
|
||||
cell = LSTMCell.from_params(*args, **kwargs)
|
||||
layer = cls(
|
||||
cell.input_size, cell.hidden_size, cell.bias, split_gates=cell.split_gates
|
||||
)
|
||||
layer.cell = cell
|
||||
return layer
|
||||
|
||||
|
||||
class _LSTMLayer(torch.nn.Module):
|
||||
r"""A single bi-directional LSTM layer."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int,
|
||||
hidden_dim: int,
|
||||
bias: bool = True,
|
||||
batch_first: bool = False,
|
||||
bidirectional: bool = False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
*,
|
||||
split_gates=False,
|
||||
) -> None:
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.batch_first = batch_first
|
||||
self.bidirectional = bidirectional
|
||||
self.layer_fw = _LSTMSingleLayer(
|
||||
input_dim, hidden_dim, bias=bias, split_gates=split_gates, **factory_kwargs
|
||||
)
|
||||
if self.bidirectional:
|
||||
self.layer_bw = _LSTMSingleLayer(
|
||||
input_dim,
|
||||
hidden_dim,
|
||||
bias=bias,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = None):
|
||||
if self.batch_first:
|
||||
x = x.transpose(0, 1)
|
||||
if hidden is None:
|
||||
hx_fw, cx_fw = (None, None)
|
||||
else:
|
||||
hx_fw, cx_fw = hidden
|
||||
hidden_bw: tuple[Tensor, Tensor] | None = None
|
||||
if self.bidirectional:
|
||||
if hx_fw is None:
|
||||
hx_bw = None
|
||||
else:
|
||||
hx_bw = hx_fw[1]
|
||||
hx_fw = hx_fw[0]
|
||||
if cx_fw is None:
|
||||
cx_bw = None
|
||||
else:
|
||||
cx_bw = cx_fw[1]
|
||||
cx_fw = cx_fw[0]
|
||||
if hx_bw is not None and cx_bw is not None:
|
||||
hidden_bw = hx_bw, cx_bw
|
||||
if hx_fw is None and cx_fw is None:
|
||||
hidden_fw = None
|
||||
else:
|
||||
hidden_fw = (
|
||||
torch.jit._unwrap_optional(hx_fw),
|
||||
torch.jit._unwrap_optional(cx_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
|
||||
c = None
|
||||
elif hidden_fw is None:
|
||||
(h, c) = torch.jit._unwrap_optional(hidden_bw)
|
||||
elif hidden_bw is None:
|
||||
(h, c) = torch.jit._unwrap_optional(hidden_fw)
|
||||
else:
|
||||
h = torch.stack([hidden_fw[0], hidden_bw[0]], 0) # type: ignore[list-item]
|
||||
c = torch.stack([hidden_fw[1], hidden_bw[1]], 0) # type: ignore[list-item]
|
||||
else:
|
||||
result = result_fw
|
||||
h, c = torch.jit._unwrap_optional(hidden_fw) # type: ignore[assignment]
|
||||
|
||||
if self.batch_first:
|
||||
result.transpose_(0, 1)
|
||||
|
||||
return result, (h, c)
|
||||
|
||||
@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.
|
||||
"""
|
||||
if not hasattr(other, "qconfig") and qconfig is None:
|
||||
raise AssertionError("other must have qconfig or qconfig must be provided")
|
||||
|
||||
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)
|
||||
split_gates = kwargs.get("split_gates", False)
|
||||
|
||||
layer = cls(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
input_size,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
hidden_size,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
bias,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
batch_first,
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
bidirectional,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
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 = _LSTMSingleLayer.from_params(
|
||||
wi, wh, bi, bh, split_gates=split_gates
|
||||
)
|
||||
|
||||
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 = _LSTMSingleLayer.from_params(
|
||||
wi, wh, bi, bh, split_gates=split_gates
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
class LSTM(torch.nn.Module):
|
||||
r"""A quantizable long short-term memory (LSTM).
|
||||
|
||||
For the description and the argument types, please, refer to :class:`~torch.nn.LSTM`
|
||||
|
||||
Attributes:
|
||||
layers : instances of the `_LSTMLayer`
|
||||
|
||||
.. note::
|
||||
To access the weights and biases, you need to access them per layer.
|
||||
See examples below.
|
||||
|
||||
Examples::
|
||||
|
||||
>>> import torch.ao.nn.quantizable as nnqa
|
||||
>>> rnn = nnqa.LSTM(10, 20, 2)
|
||||
>>> input = torch.randn(5, 3, 10)
|
||||
>>> h0 = torch.randn(2, 3, 20)
|
||||
>>> c0 = torch.randn(2, 3, 20)
|
||||
>>> output, (hn, cn) = rnn(input, (h0, c0))
|
||||
>>> # To get the weights:
|
||||
>>> # xdoctest: +SKIP
|
||||
>>> print(rnn.layers[0].weight_ih)
|
||||
tensor([[...]])
|
||||
>>> print(rnn.layers[0].weight_hh)
|
||||
AssertionError: There is no reverse path in the non-bidirectional layer
|
||||
"""
|
||||
|
||||
_FLOAT_MODULE = torch.nn.LSTM
|
||||
|
||||
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,
|
||||
*,
|
||||
split_gates: bool = False,
|
||||
) -> 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 # Default to eval mode. If we want to train, we will explicitly set to training.
|
||||
|
||||
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 LSTM is ignored. "
|
||||
"If you are training, please, use nn.LSTM version "
|
||||
"followed by `prepare` step.",
|
||||
stacklevel=2,
|
||||
)
|
||||
if num_layers == 1:
|
||||
warnings.warn(
|
||||
"dropout option adds dropout after all but last "
|
||||
"recurrent layer, so non-zero dropout expects "
|
||||
f"num_layers greater than 1, but got dropout={dropout} "
|
||||
f"and num_layers={num_layers}",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
layers = [
|
||||
_LSTMLayer(
|
||||
self.input_size,
|
||||
self.hidden_size,
|
||||
self.bias,
|
||||
batch_first=False,
|
||||
bidirectional=self.bidirectional,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
]
|
||||
layers.extend(
|
||||
_LSTMLayer(
|
||||
self.hidden_size,
|
||||
self.hidden_size,
|
||||
self.bias,
|
||||
batch_first=False,
|
||||
bidirectional=self.bidirectional,
|
||||
split_gates=split_gates,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(1, num_layers)
|
||||
)
|
||||
self.layers = torch.nn.ModuleList(layers)
|
||||
|
||||
def forward(self, x: Tensor, hidden: tuple[Tensor, Tensor] | None = 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
|
||||
)
|
||||
hxcx = [(zeros, zeros) for _ in range(self.num_layers)]
|
||||
else:
|
||||
hidden_non_opt = torch.jit._unwrap_optional(hidden)
|
||||
if isinstance(hidden_non_opt[0], Tensor):
|
||||
hx = hidden_non_opt[0].reshape(
|
||||
self.num_layers, num_directions, max_batch_size, self.hidden_size
|
||||
)
|
||||
cx = hidden_non_opt[1].reshape(
|
||||
self.num_layers, num_directions, max_batch_size, self.hidden_size
|
||||
)
|
||||
hxcx = [
|
||||
(hx[idx].squeeze(0), cx[idx].squeeze(0))
|
||||
for idx in range(self.num_layers)
|
||||
]
|
||||
else:
|
||||
hxcx = hidden_non_opt
|
||||
|
||||
hx_list = []
|
||||
cx_list = []
|
||||
for idx, layer in enumerate(self.layers):
|
||||
x, (h, c) = layer(x, hxcx[idx])
|
||||
hx_list.append(torch.jit._unwrap_optional(h))
|
||||
cx_list.append(torch.jit._unwrap_optional(c))
|
||||
hx_tensor = torch.stack(hx_list)
|
||||
cx_tensor = torch.stack(cx_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])
|
||||
cx_tensor = cx_tensor.reshape(-1, cx_tensor.shape[-2], cx_tensor.shape[-1])
|
||||
|
||||
if self.batch_first:
|
||||
x = x.transpose(0, 1)
|
||||
|
||||
return x, (hx_tensor, cx_tensor)
|
||||
|
||||
def _get_name(self):
|
||||
return "QuantizableLSTM"
|
||||
|
||||
@classmethod
|
||||
def from_float(cls, other, qconfig=None, split_gates=False):
|
||||
if not isinstance(other, cls._FLOAT_MODULE):
|
||||
raise AssertionError(
|
||||
f"Expected module type {cls._FLOAT_MODULE}, got {type(other)}"
|
||||
)
|
||||
if not hasattr(other, "qconfig") and not qconfig:
|
||||
raise AssertionError("other must have qconfig or qconfig must be provided")
|
||||
observed = cls(
|
||||
other.input_size,
|
||||
other.hidden_size,
|
||||
other.num_layers,
|
||||
other.bias,
|
||||
other.batch_first,
|
||||
other.dropout,
|
||||
other.bidirectional,
|
||||
split_gates=split_gates,
|
||||
)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
observed.qconfig = getattr(other, "qconfig", qconfig)
|
||||
for idx in range(other.num_layers):
|
||||
observed.layers[idx] = _LSTMLayer.from_float(
|
||||
other, idx, qconfig, batch_first=False, split_gates=split_gates
|
||||
)
|
||||
|
||||
# Prepare the model
|
||||
if other.training:
|
||||
observed.train()
|
||||
observed = torch.ao.quantization.prepare_qat(observed, inplace=True)
|
||||
else:
|
||||
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 LSTM module. Please, see "
|
||||
"the examples on quantizable LSTMs."
|
||||
)
|
||||
Reference in New Issue
Block a user