Improvements from: CorentinJ#126

encoder/train.py

@@ -9,6 +9,8 @@
 def sync(device: torch.device):
     # For correct profiling (cuda operations are async)
     if device.type == "cuda":
+        torch.cuda.synchronize()
+    else:
         torch.cuda.synchronize(device)
 
 
@@ -29,7 +31,7 @@ def train(run_id: str, clean_data_root: Path, models_dir: Path, umap_every: int,
     # hyperparameters) faster on the CPU.
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # FIXME: currently, the gradient is None if loss_device is cuda
-    loss_device = torch.device("cpu")
+    loss_device = torch.device("cuda")
 
     # Create the model and the optimizer
     model = SpeakerEncoder(device, loss_device)

encoder/visualizations.py

@@ -8,20 +8,40 @@
 import umap
 
 colormap = np.array([
-    [76, 255, 0],
-    [0, 127, 70],
-    [255, 0, 0],
-    [255, 217, 38],
-    [0, 135, 255],
-    [165, 0, 165],
-    [255, 167, 255],
-    [0, 255, 255],
-    [255, 96, 38],
-    [142, 76, 0],
-    [33, 0, 127],
+    [32, 25, 35],
+    [255, 255, 255],
+    [252, 255, 93],
+    [125, 252, 0],
+    [14, 196, 52],
+    [34, 140, 104],
+    [138, 216, 232],
+    [35, 91, 84],
+    [41, 189, 171],
+    [57, 152, 245],
+    [55, 41, 79],
+    [39, 125, 167],
+    [55, 80, 219],
+    [242, 32, 32],
+    [153, 25, 25],
+    [255, 203, 165],
+    [230, 143, 102],
+    [197, 97, 51],
+    [150, 52, 28],
+    [99, 40, 25],
+    [255, 196, 19],
+    [244, 122, 34],
+    [47, 42, 160],
+    [183, 50, 204],
+    [119, 43, 157],
+    [240, 124, 171],
+    [211, 11, 148],
+    [237, 239, 243],
+    [195, 165, 180],
+    [148, 106, 162],
+    [93, 76, 134],
     [0, 0, 0],
     [183, 183, 183],
-], dtype=np.float) / 255 
+], dtype=np.float) / 255
 
 
 class Visualizations:

resources/convert.sh

@@ -36,7 +36,7 @@ for f in $(find . -name "*.m4a"); do
         let skipped=skipped+1
         continue
     fi
-    run_with_lock ffmpeg -loglevel panic -i "$f" -ar 16000 "${f%.*}.wav"
+    run_with_lock ffmpeg -loglevel panic -i "$f" -ar 24000 "${f%.*}.wav"
     let converted=converted+1
     echo -ne "Total files: $total. Skipped $skipped already converted files; converted $converted new files."\\r
 done