Class to quantize given activations. Note that when using this function, the input activation quantization range will be fixed for all tokens/images for inference. This generally will affect some accuracy but achieve better latency performance. Parameters: ---------- act_range_m
| 15 | |
| 16 | |
| 17 | class QuantAct(nn.Module): |
| 18 | """ |
| 19 | Class to quantize given activations. Note that when using this function, the input activation quantization range will be fixed for all |
| 20 | tokens/images for inference. This generally will affect some accuracy but achieve better latency performance. |
| 21 | Parameters: |
| 22 | ---------- |
| 23 | act_range_momentum : float, default 0.95 |
| 24 | Momentum for updating the activation quantization range. |
| 25 | quant_mode : str, default 'symmetric' |
| 26 | """ |
| 27 | |
| 28 | def __init__(self, act_range_momentum=0.95, quant_mode='symmetric'): |
| 29 | super(QuantAct, self).__init__() |
| 30 | |
| 31 | self.act_range_momentum = act_range_momentum |
| 32 | self.quant_mode = quant_mode |
| 33 | if quant_mode == 'symmetric': |
| 34 | self.act_function = SymQuantizer.apply |
| 35 | else: |
| 36 | self.act_function = AsymQuantizer.apply |
| 37 | |
| 38 | self.register_buffer('x_min_max', torch.zeros(2)) |
| 39 | |
| 40 | def forward(self, x, num_bits, *args): |
| 41 | """ |
| 42 | x: the activation that we need to quantize |
| 43 | num_bits: the number of bits we need to quantize the activation to |
| 44 | *args: some extra arguments that are useless but needed for align with the interface of other quantization functions |
| 45 | """ |
| 46 | |
| 47 | if self.training: |
| 48 | x_min = x.data.min() |
| 49 | x_max = x.data.max() |
| 50 | |
| 51 | # Initialization |
| 52 | if self.x_min_max[0] == self.x_min_max[1]: |
| 53 | self.x_min_max[0] = x_min |
| 54 | self.x_min_max[1] = x_max |
| 55 | |
| 56 | # if do not need momentum, please set self.act_range_momentum = 0 |
| 57 | self.x_min_max[0] = self.x_min_max[0] * self.act_range_momentum + x_min * (1 - self.act_range_momentum) |
| 58 | self.x_min_max[1] = self.x_min_max[1] * self.act_range_momentum + x_max * (1 - self.act_range_momentum) |
| 59 | |
| 60 | x_q = self.act_function(x, num_bits, self.x_min_max[0], self.x_min_max[1]) |
| 61 | |
| 62 | return x_q |
| 63 | |
| 64 | |
| 65 | class Embedding_Compress(nn.Embedding): |
no outgoing calls
no test coverage detected