(input: Tensor, lower_bd: float, upper_bd: float, bitwidth: int = 8, q_type: str = "noise")
| 37 | return out_dict |
| 38 | |
| 39 | def fake_quantize_ste(input: Tensor, lower_bd: float, upper_bd: float, bitwidth: int = 8, q_type: str = "noise"): |
| 40 | q_step = (upper_bd - lower_bd) / (2**bitwidth - 1) |
| 41 | |
| 42 | if q_type == "round": |
| 43 | output_value = STE.apply(input, bitwidth, lower_bd, upper_bd) |
| 44 | elif q_type == "noise": |
| 45 | input = torch.clamp(input, lower_bd, upper_bd) |
| 46 | noise = torch.empty_like(input).uniform_(-0.5, 0.5) |
| 47 | output_value = input + noise * q_step # whether to exclude the data pts that overflows/underflows? |
| 48 | |
| 49 | out_dict = { |
| 50 | "output_value": output_value, |
| 51 | "q_step": q_step # return q_step as Q for entropy model |
| 52 | } |
| 53 | |
| 54 | return out_dict |
| 55 | |
| 56 | |
| 57 | class STE(torch.autograd.Function): |
no test coverage detected