(
self,
layer: Layer,
weight_bits: int = 8,
activation_bits: int = 8,
moving_rate: float = 0.9,
weight_quantize_type: _QuantType = 'abs_max',
activation_quantize_type: _QuantType = 'abs_max',
weight_pre_layer: Layer | None = None,
act_pre_layer: Layer | None = None,
weight_quant_layer: Layer | None = None,
act_quant_layer: Layer | None = None,
)
| 779 | name: str |
| 780 | |
| 781 | def __init__( |
| 782 | self, |
| 783 | layer: Layer, |
| 784 | weight_bits: int = 8, |
| 785 | activation_bits: int = 8, |
| 786 | moving_rate: float = 0.9, |
| 787 | weight_quantize_type: _QuantType = 'abs_max', |
| 788 | activation_quantize_type: _QuantType = 'abs_max', |
| 789 | weight_pre_layer: Layer | None = None, |
| 790 | act_pre_layer: Layer | None = None, |
| 791 | weight_quant_layer: Layer | None = None, |
| 792 | act_quant_layer: Layer | None = None, |
| 793 | ) -> None: |
| 794 | super().__init__() |
| 795 | # For Linear |
| 796 | self.weight = layer.weight |
| 797 | self.bias = layer.bias |
| 798 | self.name = layer.name |
| 799 | # For FakeQuant |
| 800 | self._linear_quant_axis = 1 |
| 801 | |
| 802 | if weight_quant_layer is not None: |
| 803 | self._fake_quant_weight = weight_quant_layer() |
| 804 | else: |
| 805 | self._fake_quant_weight = _get_fake_quant_type( |
| 806 | weight_quantize_type, |
| 807 | name=self.weight.name, |
| 808 | moving_rate=moving_rate, |
| 809 | quant_bits=weight_bits, |
| 810 | dtype=self._dtype, |
| 811 | quant_on_weight=True, |
| 812 | channel_num=self.weight.shape[self._linear_quant_axis], |
| 813 | quant_axis=self._linear_quant_axis, |
| 814 | quant_linear=True, |
| 815 | ) |
| 816 | |
| 817 | if act_quant_layer is not None: |
| 818 | self._fake_quant_input = act_quant_layer() |
| 819 | else: |
| 820 | self._fake_quant_input = _get_fake_quant_type( |
| 821 | activation_quantize_type, |
| 822 | name=layer.full_name(), |
| 823 | moving_rate=moving_rate, |
| 824 | quant_bits=activation_bits, |
| 825 | dtype=self._dtype, |
| 826 | quant_on_weight=False, |
| 827 | ) |
| 828 | |
| 829 | self._act_preprocess = ( |
| 830 | act_pre_layer() if act_pre_layer is not None else None |
| 831 | ) |
| 832 | self._weight_preprocess = ( |
| 833 | weight_pre_layer() if weight_pre_layer is not None else None |
| 834 | ) |
| 835 | |
| 836 | def forward(self, input: Tensor) -> Tensor: |
| 837 | if self._act_preprocess is not None: |
nothing calls this directly
no test coverage detected