fix

lightning_module.py

@@ -11,13 +11,6 @@ from model import (
     MultiScaleSpectralLoss,
     GSTModule,
 )
-from utils import (
-    manual_logging,
-    load_vocoder,
-    plot_and_save_mels,
-    plot_and_save_mels_all,
-)
-
 
 class PretrainLightningModule(pl.LightningModule):
     def __init__(self, config):
@@ -32,7 +25,7 @@ class PretrainLightningModule(pl.LightningModule):
             self.channelfeats = ChannelFeatureModule(config)
 
         self.channel = ChannelModule(config)
-        self.vocoder = load_vocoder(config)
+        self.vocoder = None
 
         self.criteria_a = MultiScaleSpectralLoss(config)
         if "feature_loss" in config["train"]:
@@ -154,8 +147,6 @@ class PretrainLightningModule(pl.LightningModule):
             prog_bar=True,
             logger=True,
         )
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][0], data_type="image")
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][1], data_type="audio")
 
     def test_step(self, batch, batch_idx):
         if self.config["general"]["use_gst"]:
@@ -224,24 +215,6 @@ class PretrainLightningModule(pl.LightningModule):
                         sample_rate=self.config["preprocess"]["sampling_rate"],
                         channels_first=True,
                     )
-                    plot_and_save_mels(
-                        out[key][0, ...].cpu(),
-                        mel_dir / "{}-{}.png".format(idx, key),
-                        self.config,
-                    )
-            plot_and_save_mels_all(
-                out,
-                [
-                    "reconstructed",
-                    "remastered",
-                    "channeled",
-                    "input",
-                    "input_recons",
-                    "groundtruth",
-                ],
-                mel_dir / "{}-all.png".format(idx),
-                self.config,
-            )
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(
@@ -257,21 +230,6 @@ class PretrainLightningModule(pl.LightningModule):
         }
         return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
 
-    def tflogger(self, logger_dict, data_type):
-        for lg in self.logger.experiment:
-            if type(lg).__name__ == "SummaryWriter":
-                tensorboard = lg
-        for key in logger_dict.keys():
-            manual_logging(
-                logger=tensorboard,
-                item=logger_dict[key],
-                idx=0,
-                tag=key,
-                global_step=self.global_step,
-                data_type=data_type,
-                config=self.config,
-            )
-
 
 class SSLBaseModule(pl.LightningModule):
     def __init__(self, config):
@@ -299,7 +257,7 @@ class SSLBaseModule(pl.LightningModule):
                     pre_model.channelfeats.state_dict(), strict=False
                 )
 
-        self.vocoder = load_vocoder(config)
+        self.vocoder = None
         self.criteria = self.get_loss_function(config)
 
     def training_step(self, batch, batch_idx):
@@ -405,32 +363,6 @@ class SSLBaseModule(pl.LightningModule):
                         sample_rate=self.config["preprocess"]["sampling_rate"],
                         channels_first=True,
                     )
-                    plot_and_save_mels(
-                        out[key][0, ...].cpu(),
-                        mel_dir / "{}-{}.png".format(idx, key),
-                        self.config,
-                    )
-            plot_and_save_mels_all(
-                out,
-                plot_keys,
-                mel_dir / "{}-all.png".format(idx),
-                self.config,
-            )
-
-    def tflogger(self, logger_dict, data_type):
-        for lg in self.logger.experiment:
-            if type(lg).__name__ == "SummaryWriter":
-                tensorboard = lg
-        for key in logger_dict.keys():
-            manual_logging(
-                logger=tensorboard,
-                item=logger_dict[key],
-                idx=0,
-                tag=key,
-                global_step=self.global_step,
-                data_type=data_type,
-                config=self.config,
-            )
 
 
 class SSLStepLightningModule(SSLBaseModule):
@@ -511,8 +443,6 @@ class SSLStepLightningModule(SSLBaseModule):
             prog_bar=True,
             logger=True,
         )
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][0], data_type="image")
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][1], data_type="audio")
 
     def optimizer_step(
         self,
@@ -754,8 +684,6 @@ class SSLDualLightningModule(SSLBaseModule):
             prog_bar=True,
             logger=True,
         )
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][0], data_type="image")
-        self.tflogger(logger_dict=outputs[-1]["logger_dict"][1], data_type="audio")
 
     def test_step(self, batch, batch_idx):
         if self.config["general"]["use_gst"]:
@@ -833,17 +761,6 @@ class SSLDualLightningModule(SSLBaseModule):
                         sample_rate=self.config["preprocess"]["sampling_rate"],
                         channels_first=True,
                     )
-                    plot_and_save_mels(
-                        out[key][0, ...].cpu(),
-                        mel_dir / "{}-{}.png".format(idx, key),
-                        self.config,
-                    )
-            plot_and_save_mels_all(
-                out,
-                plot_keys,
-                mel_dir / "{}-all.png".format(idx),
-                self.config,
-            )
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(

utils.py

@@ -3,18 +3,6 @@ import json
 import torch
 import torchaudio
 
-def load_vocoder(config):
-    with open(
-        "hifigan/config_{}.json".format(config["general"]["feature_type"]), "r"
-    ) as f:
-        config_hifigan = hifigan.AttrDict(json.load(f))
-    vocoder = hifigan.Generator(config_hifigan)
-    vocoder.load_state_dict(torch.load(config["general"]["hifigan_path"])["generator"])
-    vocoder.remove_weight_norm()
-    for param in vocoder.parameters():
-        param.requires_grad = False
-    return vocoder
-
 def configure_args(config, args):
     for key in ["stage", "corpus_type", "source_path", "aux_path", "preprocessed_path"]:
         if getattr(args, key) != None: