def model_size()

in diffq/diffq.py [0:0]

• 20 lines of code
• 6 McCabe index (conditional complexity)

    def model_size(self, exact=False):
        """
        Differentiable estimate of the model size.
        The size is returned in MB.

        If `exact` is True, then the output is no longer differentiable but
        reflect exactly an achievable size, even without compression,
        i.e.same as returned by `naive_model_size()`.
        """
        total = super().model_size()
        subtotal = 0
        for qparam in self._qparams:
            # only count the first appearance of a Parameter
            if qparam.other is not None:
                continue
            bits = self.extra_bits + self._get_bits(qparam.logit)
            if exact:
                bits = bits.round().clamp(1, 15)
            if self.group_size == 0:
                group_size = qparam.param.numel()
            else:
                group_size = self.group_size
            subtotal += group_size * bits.sum()
            subtotal += 2 * 32  # param scale

            # Number of bits to represent each number of bits
            bits_bits = math.ceil(math.log2(1 + (bits.max().round().item() - self.min_bits)))
            subtotal += 8  # 8 bits for bits_bits
            subtotal += bits_bits * bits.numel()

        subtotal /= 2 ** 20 * 8  # bits -> MegaBytes
        return total + subtotal