| 19 | |
| 20 | |
| 21 | class TQT_numpy: |
| 22 | def __init__(self, lowerbound, upperbound): |
| 23 | super().__init__() |
| 24 | self.lowerbound = lowerbound |
| 25 | self.upperbound = upperbound |
| 26 | |
| 27 | def forward(self, inp, scale): |
| 28 | t = 2 ** scale |
| 29 | # t = F.maximum(t, 1e-4) |
| 30 | inp_scaled = inp / t |
| 31 | inp_clipped = np.maximum( |
| 32 | np.minimum(inp_scaled, self.upperbound), self.lowerbound |
| 33 | ) |
| 34 | inp_rounded = np.round(inp_clipped) |
| 35 | inp_flq = inp_rounded * t |
| 36 | self.saved_tensors = (inp_scaled, inp_rounded, t) |
| 37 | return inp_flq |
| 38 | |
| 39 | def backward(self, grad_inp_flq): |
| 40 | (inp_scaled, inp_rounded, t) = self.saved_tensors |
| 41 | mask_clip = (inp_scaled < -0.5 + self.lowerbound) + ( |
| 42 | inp_scaled > self.upperbound + 0.5 |
| 43 | ) # mask for accumulating the gradients of |data_scaled|>L |
| 44 | mask_quant = np.abs( |
| 45 | mask_clip - 1 |
| 46 | ) # mask for accumulating the gradients with |data_scaled|<=L |
| 47 | grad_quant = ( |
| 48 | grad_inp_flq * mask_quant * (inp_rounded - inp_scaled) |
| 49 | ) # gradient within |data_scaled|<=L |
| 50 | grad_clip = ( |
| 51 | grad_inp_flq * mask_clip * inp_rounded |
| 52 | ) # gradient with | data_scaled|>L |
| 53 | grad_s = grad_clip.sum() + grad_quant.sum() |
| 54 | # dL/ds = dL/dt * t * ln(2) |
| 55 | grad_s = grad_s * t * np.log(2) |
| 56 | grad_inp = grad_inp_flq * mask_quant |
| 57 | return grad_inp, grad_s |
| 58 | |
| 59 | |
| 60 | def test_tqt(): |