| 154 | |
| 155 | |
| 156 | class LSQ_numpy: |
| 157 | def __init__(self, lowerbound, upperbound): |
| 158 | super().__init__() |
| 159 | self.lowerbound = lowerbound |
| 160 | self.upperbound = upperbound |
| 161 | |
| 162 | def forward(self, inp, scale, zero_point, grad_scale): |
| 163 | inp_scaled = inp / scale + zero_point |
| 164 | inp_clipped = np.maximum( |
| 165 | np.minimum(inp_scaled, self.upperbound), self.lowerbound |
| 166 | ) |
| 167 | inp_rounded = np.floor(inp_clipped + 0.5) |
| 168 | inp_flq = (inp_rounded - zero_point) * scale |
| 169 | self.saved_tensors = (inp_scaled, inp_rounded, scale, grad_scale) |
| 170 | return inp_flq |
| 171 | |
| 172 | def backward(self, grad_inp_flq): |
| 173 | (inp_scaled, inp_rounded, scale, grad_scale) = self.saved_tensors |
| 174 | |
| 175 | ind_small = inp_scaled < self.lowerbound |
| 176 | ind_big = inp_scaled > self.upperbound |
| 177 | ind_middle = np.logical_xor(ind_small, ind_big) |
| 178 | ind_middle = np.abs(ind_middle - 1) |
| 179 | |
| 180 | grad_s = ( |
| 181 | ind_small * self.lowerbound |
| 182 | + ind_big * self.upperbound |
| 183 | + ind_middle * (-inp_scaled + inp_rounded) |
| 184 | ) |
| 185 | grad_s = grad_s * grad_scale * grad_inp_flq |
| 186 | grad_s = grad_s.sum() |
| 187 | grad_inp = grad_inp_flq * ind_middle |
| 188 | |
| 189 | return grad_inp, grad_s |
| 190 | |
| 191 | |
| 192 | def test_lsq(): |