adding support for parallel processing in dataset generation

remfx/datasets.py

@@ -8,15 +8,16 @@ import pytorch_lightning as pl
 import random
 from tqdm import tqdm
 from pathlib import Path
-from remfx import effects
+from remfx import effects as effect_lib
 from typing import Any, List, Dict
 from torch.utils.data import Dataset, DataLoader
 from remfx.utils import select_random_chunk
+import multiprocessing
 
 
 # https://zenodo.org/record/1193957 -> VocalSet
 
-ALL_EFFECTS = effects.Pedalboard_Effects
+ALL_EFFECTS = effect_lib.Pedalboard_Effects
 # print(ALL_EFFECTS)
 
 
@@ -146,6 +147,101 @@ def locate_files(root: str, mode: str):
     return file_list
 
 
+def parallel_process_effects(
+    chunk_idx: int,
+    proc_root: str,
+    files: list,
+    chunk_size: int,
+    effects: list,
+    effects_to_keep: list,
+    num_kept_effects: tuple,
+    shuffle_kept_effects: bool,
+    effects_to_remove: list,
+    num_removed_effects: tuple,
+    shuffle_removed_effects: bool,
+    sample_rate: int,
+    target_lufs_db: float,
+):
+    chunk = None
+    random_dataset_choice = random.choice(files)
+    while chunk is None:
+        random_file_choice = random.choice(random_dataset_choice)
+        chunk = select_random_chunk(random_file_choice, chunk_size, sample_rate)
+
+    # Sum to mono
+    if chunk.shape[0] > 1:
+        chunk = chunk.sum(0, keepdim=True)
+
+    dry = chunk
+
+    # loudness normalization
+    normalize = effect_lib.LoudnessNormalize(sample_rate, target_lufs_db=target_lufs_db)
+
+    # Apply Kept Effects
+    # Shuffle effects if specified
+    if shuffle_kept_effects:
+        effect_indices = torch.randperm(len(effects_to_keep))
+    else:
+        effect_indices = torch.arange(len(effects_to_keep))
+
+    r1 = num_kept_effects[0]
+    r2 = num_kept_effects[1]
+    num_kept_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
+    effect_indices = effect_indices[:num_kept_effects]
+    # Index in effect settings
+    effect_names_to_apply = [effects_to_keep[i] for i in effect_indices]
+    effects_to_apply = [effects[i] for i in effect_names_to_apply]
+    # Apply
+    dry_labels = []
+    for effect in effects_to_apply:
+        # Normalize in-between effects
+        dry = normalize(effect(dry))
+        dry_labels.append(ALL_EFFECTS.index(type(effect)))
+
+    # Apply effects_to_remove
+    # Shuffle effects if specified
+    if shuffle_removed_effects:
+        effect_indices = torch.randperm(len(effects_to_remove))
+    else:
+        effect_indices = torch.arange(len(effects_to_remove))
+    wet = torch.clone(dry)
+    r1 = num_removed_effects[0]
+    r2 = num_removed_effects[1]
+    num_removed_effects = torch.round((r1 - r2) * torch.rand(1) + r2).int()
+    effect_indices = effect_indices[:num_removed_effects]
+    # Index in effect settings
+    effect_names_to_apply = [effects_to_remove[i] for i in effect_indices]
+    effects_to_apply = [effects[i] for i in effect_names_to_apply]
+    # Apply
+    wet_labels = []
+    for effect in effects_to_apply:
+        # Normalize in-between effects
+        wet = normalize(effect(wet))
+        wet_labels.append(ALL_EFFECTS.index(type(effect)))
+
+    wet_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+    dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
+
+    for label_idx in wet_labels:
+        wet_labels_tensor[label_idx] = 1.0
+
+    for label_idx in dry_labels:
+        dry_labels_tensor[label_idx] = 1.0
+
+    # Normalize
+    normalized_dry = normalize(dry)
+    normalized_wet = normalize(wet)
+
+    output_dir = proc_root / str(chunk_idx)
+    output_dir.mkdir(exist_ok=True)
+    torchaudio.save(output_dir / "input.wav", normalized_wet, sample_rate)
+    torchaudio.save(output_dir / "target.wav", normalized_dry, sample_rate)
+    torch.save(dry_labels_tensor, output_dir / "dry_effects.pt")
+    torch.save(wet_labels_tensor, output_dir / "wet_effects.pt")
+
+    # return normalized_dry, normalized_wet, dry_labels_tensor, wet_labels_tensor
+
+
 class EffectDataset(Dataset):
     def __init__(
         self,
@@ -163,6 +259,7 @@ class EffectDataset(Dataset):
         render_files: bool = True,
         render_root: str = None,
         mode: str = "train",
+        parallel: bool = True,
     ):
         super().__init__()
         self.chunks = []
@@ -177,7 +274,7 @@ class EffectDataset(Dataset):
         self.num_removed_effects = num_removed_effects
         self.effects_to_keep = [] if effects_to_keep is None else effects_to_keep
         self.effects_to_remove = [] if effects_to_remove is None else effects_to_remove
-        self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
+        self.normalize = effect_lib.LoudnessNormalize(sample_rate, target_lufs_db=-20)
         self.effects = effect_modules
         self.shuffle_kept_effects = shuffle_kept_effects
         self.shuffle_removed_effects = shuffle_removed_effects
@@ -192,6 +289,7 @@ class EffectDataset(Dataset):
         )
         self.validate_effect_input()
         self.proc_root = self.render_root / "processed" / effects_string / self.mode
+        self.parallel = parallel
 
         self.files = locate_files(self.root, self.mode)
 
@@ -212,26 +310,50 @@ class EffectDataset(Dataset):
         if render_files:
             # Split audio file into chunks, resample, then apply random effects
             self.proc_root.mkdir(parents=True, exist_ok=True)
-            for num_chunk in tqdm(range(self.total_chunks)):
-                chunk = None
-                random_dataset_choice = random.choice(self.files)
-                while chunk is None:
-                    random_file_choice = random.choice(random_dataset_choice)
-                    chunk = select_random_chunk(
-                        random_file_choice, self.chunk_size, self.sample_rate
-                    )
 
-                # Sum to mono
-                if chunk.shape[0] > 1:
-                    chunk = chunk.sum(0, keepdim=True)
-
-                dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
-                output_dir = self.proc_root / str(num_chunk)
-                output_dir.mkdir(exist_ok=True)
-                torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
-                torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
-                torch.save(dry_effects, output_dir / "dry_effects.pt")
-                torch.save(wet_effects, output_dir / "wet_effects.pt")
+            if self.parallel:
+                items = [
+                    (
+                        chunk_idx,
+                        self.proc_root,
+                        self.files,
+                        self.chunk_size,
+                        self.effects,
+                        self.effects_to_keep,
+                        self.num_kept_effects,
+                        self.shuffle_kept_effects,
+                        self.effects_to_remove,
+                        self.num_removed_effects,
+                        self.shuffle_removed_effects,
+                        self.sample_rate,
+                        -20.0,
+                    )
+                    for chunk_idx in range(self.total_chunks)
+                ]
+                with multiprocessing.Pool(processes=32) as pool:
+                    pool.starmap(parallel_process_effects, items)
+                print(f"Done proccessing {self.total_chunks}", flush=True)
+            else:
+                for num_chunk in tqdm(range(self.total_chunks)):
+                    chunk = None
+                    random_dataset_choice = random.choice(self.files)
+                    while chunk is None:
+                        random_file_choice = random.choice(random_dataset_choice)
+                        chunk = select_random_chunk(
+                            random_file_choice, self.chunk_size, self.sample_rate
+                        )
+
+                    # Sum to mono
+                    if chunk.shape[0] > 1:
+                        chunk = chunk.sum(0, keepdim=True)
+
+                    dry, wet, dry_effects, wet_effects = self.process_effects(chunk)
+                    output_dir = self.proc_root / str(num_chunk)
+                    output_dir.mkdir(exist_ok=True)
+                    torchaudio.save(output_dir / "input.wav", wet, self.sample_rate)
+                    torchaudio.save(output_dir / "target.wav", dry, self.sample_rate)
+                    torch.save(dry_effects, output_dir / "dry_effects.pt")
+                    torch.save(wet_effects, output_dir / "wet_effects.pt")
 
             print("Finished rendering")
         else: