Add gradient clipping and lr scheduler

cfg/config.yaml

@@ -8,7 +8,6 @@ train: True
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-log_every_n_steps: 1000
 render_files: True
 render_root: "./data/processed"
 
@@ -22,6 +21,9 @@ callbacks:
     verbose: False
     dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
     filename: '{epoch:02d}-{valid_loss:.3f}'
+  learning_rate_monitor:
+    _target_: pytorch_lightning.callbacks.LearningRateMonitor
+    logging_interval: "step"
 
 datamodule:
   _target_: remfx.datasets.VocalSetDatamodule
@@ -77,3 +79,5 @@ trainer:
   accumulate_grad_batches: 1
   accelerator: null
   devices: 1
+  gradient_clip_val: 10.0
+  max_steps: 50000

remfx/datasets.py

@@ -17,7 +17,7 @@ class VocalSet(Dataset):
         self,
         root: str,
         sample_rate: int,
-        chunk_size_in_sec: int = 3,
+        chunk_size: int = 3,
         effect_types: List[torch.nn.Module] = None,
         render_files: bool = True,
         render_root: str = None,
@@ -28,7 +28,7 @@ class VocalSet(Dataset):
         self.song_idx = []
         self.root = Path(root)
         self.render_root = Path(render_root)
-        self.chunk_size_in_sec = chunk_size_in_sec
+        self.chunk_size = chunk_size
         self.sample_rate = sample_rate
         self.mode = mode
 
@@ -46,15 +46,12 @@ class VocalSet(Dataset):
             # Split audio file into chunks, resample, then apply random effects
             self.processed_root.mkdir(parents=True, exist_ok=True)
             for audio_file in tqdm(self.files, total=len(self.files)):
-                chunks, orig_sr = create_sequential_chunks(
-                    audio_file, self.chunk_size_in_sec
-                )
+                chunks, orig_sr = create_sequential_chunks(audio_file, self.chunk_size)
                 for chunk in chunks:
                     resampled_chunk = torchaudio.functional.resample(
                         chunk, orig_sr, sample_rate
                     )
-                    chunk_size_in_samples = self.chunk_size_in_sec * self.sample_rate
-                    if resampled_chunk.shape[-1] < chunk_size_in_samples:
+                    if resampled_chunk.shape[-1] < chunk_size:
                         # Skip if chunk is too small
                         continue
                     # Apply effect

remfx/models.py

@@ -55,6 +55,28 @@ class RemFXModel(pl.LightningModule):
         )
         return optimizer
 
+    # Add step-based learning rate scheduler
+    def optimizer_step(
+        self,
+        epoch,
+        batch_idx,
+        optimizer,
+        optimizer_idx,
+        optimizer_closure,
+        on_tpu=False,
+        using_lbfgs=False,
+    ):
+        # update params
+        optimizer.step(closure=optimizer_closure)
+
+        # update learning rate. Reduce by factor of 10 at 80% and 95% of training
+        if self.trainer.global_step == 0.8 * self.trainer.max_steps:
+            for pg in optimizer.param_groups:
+                pg["lr"] = 0.1 * pg["lr"]
+        if self.trainer.global_step == 0.95 * self.trainer.max_steps:
+            for pg in optimizer.param_groups:
+                pg["lr"] = 0.1 * pg["lr"]
+
     def training_step(self, batch, batch_idx):
         loss = self.common_step(batch, batch_idx, mode="train")
         return loss

remfx/utils.py

@@ -132,10 +132,9 @@ def create_sequential_chunks(
     """
     chunks = []
     audio, sr = torchaudio.load(audio_file)
-    chunk_size_in_samples = chunk_size * sr
-    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size_in_samples)
+    chunk_starts = torch.arange(0, audio.shape[-1], chunk_size)
     for start in chunk_starts:
-        if start + chunk_size_in_samples > audio.shape[-1]:
+        if start + chunk_size > audio.shape[-1]:
             break
-        chunks.append(audio[:, start : start + chunk_size_in_samples])
+        chunks.append(audio[:, start : start + chunk_size])
     return chunks, sr

scripts/train.py

@@ -42,6 +42,7 @@ def main(cfg: DictConfig):
     summary = ModelSummary(model)
     print(summary)
     trainer.fit(model=model, datamodule=datamodule)
+    trainer.test(model=model, datamodule=datamodule)
 
 
 if __name__ == "__main__":