def on_start()

in classy_vision/hooks/model_complexity_hook.py [0:0]

• 41 lines of code
• 7 McCabe index (conditional complexity)

    def on_start(self, task) -> None:
        """Measure number of parameters, FLOPs and activations."""
        self.num_flops = 0
        self.num_activations = 0
        self.num_parameters = 0
        try:
            self.num_parameters = count_params(task.base_model)
            logging.info("Number of parameters in model: %d" % self.num_parameters)
            try:
                self.num_flops = compute_flops(
                    task.base_model,
                    input_shape=task.base_model.input_shape,
                    input_key=task.base_model.input_key
                    if hasattr(task.base_model, "input_key")
                    else None,
                )
                if self.num_flops is None:
                    logging.info("FLOPs for forward pass: skipped.")
                    self.num_flops = 0
                else:
                    logging.info(
                        "FLOPs for forward pass: %d MFLOPs"
                        % (float(self.num_flops) / 1e6)
                    )
            except ClassyProfilerNotImplementedError as e:
                logging.warning(f"Could not compute FLOPs for model forward pass: {e}")
            try:
                self.num_activations = compute_activations(
                    task.base_model,
                    input_shape=task.base_model.input_shape,
                    input_key=task.base_model.input_key
                    if hasattr(task.base_model, "input_key")
                    else None,
                )
                logging.info(f"Number of activations in model: {self.num_activations}")
            except ClassyProfilerNotImplementedError as e:
                logging.warning(
                    f"Could not compute activations for model forward pass: {e}"
                )
        except Exception:
            logging.info("Skipping complexity calculation: Unexpected error")
            logging.debug("Error trace for complexity calculation:", exc_info=True)