From 090bb14344442b2f74805edbd1fc0408fc741ed8 Mon Sep 17 00:00:00 2001 From: Nick Fraser Date: Fri, 4 Oct 2024 13:59:46 +0100 Subject: [PATCH] Fix (nn/TruncAvgPool): Remove any quant tensor manual manipulation. --- src/brevitas/nn/quant_avg_pool.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/brevitas/nn/quant_avg_pool.py b/src/brevitas/nn/quant_avg_pool.py index 7a3f108da..31bac921b 100644 --- a/src/brevitas/nn/quant_avg_pool.py +++ b/src/brevitas/nn/quant_avg_pool.py @@ -54,6 +54,7 @@ def _avg_scaling(self): else: return self.kernel_size * self.kernel_size + # TODO: Replace with functional call def forward(self, input: Union[Tensor, QuantTensor]): x = self.unpack_input(input) @@ -62,8 +63,6 @@ def forward(self, input: Union[Tensor, QuantTensor]): if isinstance(x, QuantTensor) and self.is_trunc_quant_enabled: y = AvgPool2d.forward(self, x) - rescaled_value = y.value * self._avg_scaling - y = y.set(value=rescaled_value) y = self.trunc_quant(y) else: y = AvgPool2d.forward(self, _unpack_quant_tensor(x)) @@ -111,6 +110,7 @@ def compute_kernel_size_stride(input_shape, output_shape): stride_list.append(stride) return kernel_size_list, stride_list + # TODO: Replace with functional call def forward(self, input: Union[Tensor, QuantTensor]): x = self.unpack_input(input) @@ -122,10 +122,6 @@ def forward(self, input: Union[Tensor, QuantTensor]): if isinstance(x, QuantTensor) and self.is_trunc_quant_enabled: y = AdaptiveAvgPool2d.forward(self, x) - k_size, stride = self.compute_kernel_size_stride(x.value.shape[2:], y.value.shape[2:]) - reduce_size = reduce(mul, k_size, 1) - rescaled_value = y.value * reduce_size # remove avg scaling - y = y.set(value=rescaled_value) y = self.trunc_quant(y) else: y = AdaptiveAvgPool2d.forward(self, _unpack_quant_tensor(x))