Method one_step_equalization

ppq/quantization/optim/ssd.py:264–322

Equalization step with scale being calculated in the way specified by algo_type. Args: pair (List[Operation]): a list of operations representing a equalzation pair op_act_channel_range (Dict[Operation, torch.Tensor]): channel-wise activation range of all Conv

(
        self,
        pair: List[Operation],
        op_act_channel_range: Dict[Operation, torch.Tensor]={},
        algo_type: int=2,
        ssd_min_scale: float=8,
        ssd_max_scale: float=2,
        dfq_min_scale: float=0.1,
        dfq_max_scale: float=10,
        eps: float=1e-8
    )

            pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(-1, 1, 1, 1)

    def one_step_equalization(
        self,
        pair: List[Operation],
        op_act_channel_range: Dict[Operation, torch.Tensor]={},
        algo_type: int=2,
        ssd_min_scale: float=8,
        ssd_max_scale: float=2,
        dfq_min_scale: float=0.1,
        dfq_max_scale: float=10,
        eps: float=1e-8
    ):
        """Equalization step with scale being calculated in the way specified
        by algo_type.

        Args:
            pair (List[Operation]): a list of operations representing a equalzation pair
            op_act_channel_range (Dict[Operation, torch.Tensor]): channel-wise activation range of all Conv ops in the graph
            algo_type (int, optional): minor algo type. 0 represents dfq algo, 1~3 represents ssd algo. Defaults to 2.
            ssd_min_scale (float, optional): minimum clip value of scale for ssd algo. Defaults to 8.
            ssd_max_scale (float, optional): maximum clip value of scale for ssd algo. Defaults to 2.
            dfq_min_scale (float, optional): minimum clip value of scale for dfq algo. Defaults to 0.1.
            dfq_max_scale (float, optional): maximum clip value of scale for dfq algo. Defaults to 10.
            eps (float, optional): small constant for numerical stability. Defaults to 1e-8.
        """
        first_weight_range, last_weight_range = self.prepare_weight_for_equalization(pair)

        if algo_type == 0:
            # dfq
            scale = torch.sqrt(last_weight_range / (first_weight_range + eps))
            scale = torch.clamp(scale, dfq_min_scale, dfq_max_scale)

        else:
            first_weight_range = torch.where(first_weight_range < first_weight_range.max() * self.channel_ratio,\
                                        first_weight_range.max() * self.channel_ratio, first_weight_range)
            last_weight_range = torch.where(last_weight_range < last_weight_range.max() * self.channel_ratio,\
                                        last_weight_range.max() * self.channel_ratio, last_weight_range)

            kernel_scale = first_weight_range.max() / (first_weight_range + eps)
            next_kernel_scale = last_weight_range.max() / (last_weight_range + eps)
            first_weight_act_range = op_act_channel_range[pair[0]]
            first_weight_act_range = torch.where(first_weight_act_range < 0.01, torch.tensor(0.01,\
                device=first_weight_act_range.device, dtype=torch.float32), first_weight_act_range)
            act_scale = first_weight_act_range.max() / (first_weight_act_range + eps)

            if algo_type == 1:
                scale = torch.min(kernel_scale, act_scale)
            elif algo_type == 2:
                kernel_scale = kernel_scale / next_kernel_scale
                act_scale = act_scale / next_kernel_scale
                scale = torch.min(kernel_scale, act_scale)
                scale = torch.min(scale, torch.tensor(ssd_min_scale, dtype=torch.float32, device=scale.device))
                scale /= scale.min()
                scale = torch.clamp(scale, 1.0, ssd_max_scale)
            else:
                kernel_scale = (kernel_scale / next_kernel_scale).sqrt()
                scale = (act_scale * kernel_scale).sqrt()
                scale = torch.clamp(scale, 1.0, ssd_max_scale)