module owlite.nn.functions.ste
class ScaledRoundSTE
A round function that uses STE backward.
The input is divided by the scale, rounded up, and multiplied by the scale again. No gradient is propagated through the scale.
$$ \text{output} = \lfloor \text{input} / \text{scale} \rceil * \text{scale} $$
method backward
python
backward(ctx: Any, *grad_outputs: Any) → Any
method forward
python
forward(ctx: Any, inputs: Tensor, scale: Tensor | float = 1.0) → Any
class FakeQuantizeSTEFunction
Fake quantizing function for QAT using STE (Straight-Through Estimator).
For $$ quant_min $$ <= input
<= $$ quant_max $$ the gradient passes straight through, otherwise the gradient is zero
In STE(Straight Through Estimation) method, the gradient of the round function used in fake quantization is approximated as 1, and backpropagation is performed based on this approximation. As a result, the gradient of the input entering the fake quantizer is propagated as is when it falls between $$ quant_min $$ and $$ quant_max $$, while gradients outside this range become 0. However, since the gradient propagated to $$ step_size $$ is 0, $$ step_size $$ is fixed.
When $$x$$ is input of FakeQuantize .
$$ \hat{x} = \text{FakeQuantize}(x) $$
method backward
python
backward(ctx: Any, *grad_outputs: Any) → Any
method forward
python
forward(
ctx: Any,
inputs: Tensor,
step_size: Tensor,
zero_point: Tensor,
grad_scale: float,
quant_min: int,
quant_max: int,
axis: int | None
) → Tensor
Updated: 2024-06-13T23:42:41