Method write_back

ppq/quantization/optim/ssd.py:212–262

(self, pair: List[Operation], scale: torch.Tensor)

        return first_weight_range, last_weight_range

    def write_back(self, pair: List[Operation], scale: torch.Tensor) -> None:
        first_computing_op_weight = pair[0].parameters[0].value
        last_computing_op_weight  = pair[-1].parameters[0].value

        assert isinstance(first_computing_op_weight, torch.Tensor)
        assert isinstance(last_computing_op_weight, torch.Tensor)

        if pair[0].type == 'Conv':
            pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(-1, 1, 1, 1)

        elif pair[0].type == 'Gemm':
            if pair[0].attributes.get('transB', 0):
                pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(-1, 1)
            else:
                pair[0].parameters[0].value = first_computing_op_weight * scale.reshape(1, -1)

        elif pair[0].type == 'ConvTranspose':
            num_group = pair[0].attributes.get('group', 1)
            C_in, C_out_g, K1, K2 = first_computing_op_weight.shape
            first_computing_op_weight = first_computing_op_weight.reshape(num_group, C_in // num_group, C_out_g, K1, K2)
            first_computing_op_weight = first_computing_op_weight * scale.reshape(num_group, 1, -1, 1, 1)
            pair[0].parameters[0].value = first_computing_op_weight.reshape(C_in, C_out_g, K1, K2)

        if len(pair[0].parameters) > 1:
            pair[0].parameters[1].value = pair[0].parameters[1].value * scale

        if pair[-1].type == 'Conv':
            num_group = pair[-1].attributes.get('group', 1)
            C_out, C_in_g, K1, K2 = last_computing_op_weight.shape
            last_computing_op_weight = last_computing_op_weight.reshape(num_group, C_out // num_group, C_in_g, K1, K2)
            last_computing_op_weight = last_computing_op_weight / scale.reshape(num_group, 1, -1, 1, 1)
            pair[-1].parameters[0].value = last_computing_op_weight.reshape(C_out, C_in_g, K1, K2)

        elif pair[-1].type == 'Gemm':
            if pair[-1].attributes.get('transB', 0):
                if scale.numel() != last_computing_op_weight.shape[1]:
                    last_computing_op_weight = last_computing_op_weight.reshape(last_computing_op_weight.shape[0], scale.numel(), -1)
                    last_computing_op_weight = last_computing_op_weight / scale.reshape(1, -1, 1)
                    pair[-1].parameters[0].value = last_computing_op_weight.reshape(last_computing_op_weight.shape[0], -1)
                else:
                    pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(1, -1)
            else:
                if scale.numel() != last_computing_op_weight.shape[0]:
                    last_computing_op_weight = last_computing_op_weight.reshape(scale.numel(), -1, last_computing_op_weight.shape[-1])
                    last_computing_op_weight = last_computing_op_weight / scale.reshape(-1, 1, 1)
                    pair[-1].parameters[0].value = last_computing_op_weight.reshape(-1, last_computing_op_weight.shape[-1])
                else:
                    pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(-1, 1)

        elif pair[-1].type == 'ConvTranspose':
            pair[-1].parameters[0].value = last_computing_op_weight / scale.reshape(-1, 1, 1, 1)

    def one_step_equalization(
        self,