def collect_parameters()

in neuralcompression/models/scale_hyperprior.py [0:0]

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

    def collect_parameters(self) -> Tuple[Dict[str, Parameter], Dict[str, Parameter]]:
        """
        Separates the trainable parameters of the model into groups.

        The module's parameters are organized into "model parameters" (the
        parameters that dictate the function of the model) and "quantile
        parameters" (which are only used to learn the quantiles of the
        hyper_bottleneck layer's factorized distribution, for use at inference
        time).

        Returns:
            tuple of (model parameter_dict, quantile parameter_dict)
        """
        model_parameters = {
            n: p
            for n, p in self.named_parameters()
            if not n.endswith(".quantiles") and p.requires_grad
        }
        quantile_parameters = {
            n: p
            for n, p in self.named_parameters()
            if n.endswith(".quantiles") and p.requires_grad
        }

        model_keys = set(model_parameters.keys())
        quantile_keys = set(quantile_parameters.keys())

        # Make sure we don't have an intersection of parameters
        params_dict = dict(self.named_parameters())
        params_keys = set(params_dict.keys())

        inter_keys = model_keys.intersection(quantile_keys)
        union_keys = model_keys.union(quantile_keys)

        if len(inter_keys) != 0 or union_keys != params_keys:
            raise RuntimeError("Separating model and quantile parameters failed.")

        return model_parameters, quantile_parameters