def _log_activations()

in vision/m4/training/trainer.py [0:0]

• 22 lines of code
• 15 McCabe index (conditional complexity)

    def _log_activations(self, curr_opt_step):
        if not self.activation_tracker.jsonl_stats:
            return

        if LoggingTypes.JSONL in self.hparams.train_logging_activations:
            if self.hparams.job_id is not None:
                log_activations_filename = (
                    self.hparams.save_dir / "logs" / f"{self.hparams.job_id}_logs_activations.jsonl"
                )
            else:
                log_activations_filename = self.hparams.save_dir / "logs" / "logs_activations.jsonl"

            self.activation_tracker.dump_stats(log_activations_filename, curr_opt_step=curr_opt_step)

        # if LoggingTypes.WANDB in self.hparams.train_logging_activations and self.hparams.wandb_enable:
        #     for stats in self.activation_tracker.jsonl_stats:
        #         self.accelerator.log({**stats, **self._get_additional_step_logs()}, step=curr_opt_step)

        if LoggingTypes.PRINT in self.hparams.train_logging_activations:
            for stats in self.activation_tracker.jsonl_stats:
                stats["step"] = curr_opt_step
                activation_format = {
                    k: "" if ("nonzero" in k or "step" in k or "name" in k or "type" in k or "batches" in k) else "e"
                    for k in stats.keys()
                }
                log = "Activation stats: "
                log += self.format_print_logs(stats, activation_format)
                print(log)

        self.activation_tracker.reset_jsonl_stats()