Repository: coqui-ai/TTS
File: TTS/trainer.py (Python)

Code before:
        loader_start_time = time.time()
        batch = None
        for cur_step, batch in enumerate(self.eval_loader):
            # format data
            batch = self.format_batch(batch)
            loader_time = time.time() - loader_start_time
            self.keep_avg_eval.update_values({"avg_loader_time": loader_time})
            outputs, _ = self.eval_step(batch, cur_step)
        # plot epoch stats, artifacts and figures
        if self.args.rank == 0:
            figures, audios = None, None
            if hasattr(self.model, "module") and hasattr(self.model.module, "eval_log"):
                figures, audios = self.model.module.eval_log(self.ap, batch, outputs)
            elif hasattr(self.model, "eval_log"):
                figures, audios = self.model.eval_log(self.ap, batch, outputs)
            if figures is not None:
                self.tb_logger.tb_eval_figures(self.total_steps_done, figures)
            if audios is not None:
                self.tb_logger.tb_eval_audios(self.total_steps_done, audios, self.ap.sample_rate)
            self.tb_logger.tb_eval_stats(self.total_steps_done, self.keep_avg_eval.avg_values)

    def test_run(self) -> None:
        """Run test and log the results. Test run must be defined by the model.
        Model must return figures and audios to be logged by the Tensorboard."""
        if hasattr(self.model, "test_run"):
            if isinstance(self.model, Wavegrad):
                return None # TODO: Fix inference on WaveGrad
            elif hasattr(self.eval_loader.dataset, "load_test_samples"):
                samples = self.eval_loader.dataset.load_test_samples(1)
                figures, audios = self.model.test_run(self.ap, samples, None, self.use_cuda)
            else:
                figures, audios = self.model.test_run(self.ap, self.use_cuda)
            self.tb_logger.tb_test_audios(self.total_steps_done, audios, self.config.audio["sample_rate"])
            self.tb_logger.tb_test_figures(self.total_steps_done, figures)

    def _fit(self) -> None:
        """🏃 train -> evaluate -> test for the number of epochs."""
        if self.restore_step != 0 or self.args.best_path:
            print(" > Restoring best loss from " f"{os.path.basename(self.args.best_path)} ...")
            self.best_loss = torch.load(self.args.best_path, map_location="cpu")["model_loss"]
            print(f" > Starting with loaded last best loss {self.best_loss}.")

        self.total_steps_done = self.restore_step

        for epoch in range(0, self.config.epochs):
            self.callbacks.on_epoch_start()
            self.keep_avg_train = KeepAverage()
            self.keep_avg_eval = KeepAverage() if self.config.run_eval else None
            self.epochs_done = epoch
            self.c_logger.print_epoch_start(epoch, self.config.epochs)
            self.train_epoch()

Review comment (suggestion):
Here I had to set an exception for Wavegrad, it seems that the problem comes from the inference function that eats all the memory available. Not linked to this PR (and I don't know how to fix it) 

Code after:
        loader_start_time = time.time()
        batch = None
        for cur_step, batch in enumerate(self.eval_loader):
            # format data
            batch = self.format_batch(batch)
            loader_time = time.time() - loader_start_time
            self.keep_avg_eval.update_values({"avg_loader_time": loader_time})
            outputs, _ = self.eval_step(batch, cur_step)
        # plot epoch stats, artifacts and figures
        if self.args.rank == 0:
            figures, audios = None, None
            if hasattr(self.model, "module") and hasattr(self.model.module, "eval_log"):
                figures, audios = self.model.module.eval_log(self.ap, batch, outputs)
            elif hasattr(self.model, "eval_log"):
                figures, audios = self.model.eval_log(self.ap, batch, outputs)
            if figures is not None:
                self.tb_logger.tb_eval_figures(self.total_steps_done, figures)
            if audios is not None:
                self.tb_logger.tb_eval_audios(self.total_steps_done, audios, self.ap.sample_rate)
            self.tb_logger.tb_eval_stats(self.total_steps_done, self.keep_avg_eval.avg_values)

    def test_run(self) -> None:
        """Run test and log the results. Test run must be defined by the model.
        Model must return figures and audios to be logged by the Tensorboard."""
        if hasattr(self.model, "test_run"):
            if isinstance(self.model, Wavegrad):
                return None  # TODO: Fix inference on WaveGrad
            if hasattr(self.eval_loader.dataset, "load_test_samples"):
                samples = self.eval_loader.dataset.load_test_samples(1)
                figures, audios = self.model.test_run(self.ap, samples, None)
            else:
                figures, audios = self.model.test_run(self.ap)
            self.tb_logger.tb_test_audios(self.total_steps_done, audios, self.config.audio["sample_rate"])
            self.tb_logger.tb_test_figures(self.total_steps_done, figures)
        return None

    def _fit(self) -> None:
        """🏃 train -> evaluate -> test for the number of epochs."""
        if self.restore_step != 0 or self.args.best_path:
            print(" > Restoring best loss from " f"{os.path.basename(self.args.best_path)} ...")
            self.best_loss = torch.load(self.args.best_path, map_location="cpu")["model_loss"]
            print(f" > Starting with loaded last best loss {self.best_loss}.")

        self.total_steps_done = self.restore_step

        for epoch in range(0, self.config.epochs):
            self.callbacks.on_epoch_start()
            self.keep_avg_train = KeepAverage()
            self.keep_avg_eval = KeepAverage() if self.config.run_eval else None
            self.epochs_done = epoch
            self.c_logger.print_epoch_start(epoch, self.config.epochs)