Fix random crop length

README.md

@@ -3,7 +3,8 @@
 1. `python3 -m venv env`
 2. `source env/bin/activate`
 3. `pip install -e .`
-4. `pip install -e umx`
+4. `git submodule update --init --recursive`
+5. `pip install -e umx`
 
 ## Download [GuitarFX Dataset] (https://zenodo.org/record/7044411/)
 `./scripts/download_egfx.sh`

config.yaml

@@ -26,12 +26,13 @@ datamodule:
     sample_rate: ${sample_rate}
     root: ${oc.env:DATASET_ROOT}
     length: ${length}
-    chunk_size_in_sec: 3
+    chunk_size_in_sec: 6
     num_chunks: 10
   val_split: 0.2
   batch_size: 16
   num_workers: 8
   pin_memory: True
+  persistent_workers: True
 
 logger:
   _target_: pytorch_lightning.loggers.WandbLogger

remfx/datasets.py

@@ -88,6 +88,8 @@ def create_random_chunks(
     """
     audio, sr = torchaudio.load(audio_file)
     chunk_size_in_samples = chunk_size * sr
+    if chunk_size_in_samples >= audio.shape[-1]:
+        chunk_size_in_samples = audio.shape[-1] - 1
     chunks = []
     for i in range(num_chunks):
         start = torch.randint(0, audio.shape[-1] - chunk_size_in_samples, (1,)).item()

remfx/models.py

@@ -73,6 +73,7 @@ class RemFXModel(pl.LightningModule):
                 on_epoch=True,
                 logger=True,
                 prog_bar=True,
+                sync_dist=True,
             )
 
         return loss
@@ -83,18 +84,21 @@ class RemFXModel(pl.LightningModule):
     def on_validation_batch_start(self, batch, batch_idx, dataloader_idx):
         if self.log_next:
             x, target, label = batch
-            y = self.model.sample(x)
+            self.model.eval()
+            with torch.no_grad():
+                y = self.model.sample(x)
 
             # Concat samples together for easier viewing in dashboard
-            concat_samples = torch.cat([x, y, target], dim=-1)
+            concat_samples = torch.cat([y, x, target], dim=-1)
             log_wandb_audio_batch(
                 logger=self.logger,
-                id="prediction_sample_target",
+                id="prediction_input_target",
                 samples=concat_samples.cpu(),
                 sampling_rate=self.sample_rate,
                 caption=f"Epoch {self.current_epoch}",
             )
             self.log_next = False
+            self.model.train()
 
 
 class OpenUnmixModel(torch.nn.Module):

scripts/download_egfx.sh

@@ -17,5 +17,6 @@ wget https://zenodo.org/record/7044411/files/Sweep-Echo.zip?download=1 -O Sweep-
 wget https://zenodo.org/record/7044411/files/TapeEcho.zip?download=1 -O TapeEcho.zip
 wget https://zenodo.org/record/7044411/files/TubeScreamer.zip?download=1 -O TubeScreamer.zip
 unzip -n \*.zip
+rm -rf *.zip