updated paths to work remotely

models/config/train.yaml

@@ -1,6 +1,6 @@
 global:
   id: ast_ptl
-  device: mps
+  device: cuda
   seed: 42
   dance_ids:
     - ATN
@@ -19,10 +19,10 @@ global:
     - VWZ
     - WCS
 data_module:
-  song_data_path: data/samples/songs_cleaned.csv
-  song_audio_path: data/samples
-  batch_size: 256
-  num_workers: 10
+  song_data_path: ../datastores/dance-music/songs_cleaned.csv
+  song_audio_path: ../datastores/dance-music
+  batch_size: 64
+  num_workers: 4
   min_votes: 1
   dataset_kwargs:
     audio_window_duration: 6

preprocessing/pipelines.py

@@ -3,8 +3,6 @@ import torchaudio
 from torchaudio import transforms as taT, functional as taF
 import torch.nn as nn
 
-NOISE_PATH = "data/augmentation/Lab41-SRI-VOiCES-rm1-babb-mc01-stu-clo.wav"
-
 class AudioTrainingPipeline(torch.nn.Module):
     def __init__(self, 
             input_freq=16000,
@@ -13,12 +11,13 @@ class AudioTrainingPipeline(torch.nn.Module):
             freq_mask_size=10,
             time_mask_size=80,
             mask_count = 2,
-            snr_mean=6.0):
+            snr_mean=6.0,
+            noise_path=None):
         super().__init__()
         self.input_freq = input_freq
         self.snr_mean = snr_mean
         self.mask_count = mask_count
-        self.noise = self.get_noise()
+        self.noise = self.get_noise(noise_path)
         self.resample = taT.Resample(input_freq,resample_freq)
         self.preprocess_waveform = WaveformPreprocessing(resample_freq * expected_duration)
         self.audio_to_spectrogram = AudioToSpectrogram(
@@ -28,8 +27,10 @@ class AudioTrainingPipeline(torch.nn.Module):
         self.time_mask = taT.TimeMasking(time_mask_size)
         
 
-    def get_noise(self) -> torch.Tensor:
-        noise, sr = torchaudio.load(NOISE_PATH)
+    def get_noise(self, path) -> torch.Tensor:
+        if path is None:
+            return None
+        noise, sr = torchaudio.load(path)
         if noise.shape[0] > 1:
             noise = noise.mean(0, keepdim=True)
         if sr != self.input_freq:
@@ -37,6 +38,7 @@ class AudioTrainingPipeline(torch.nn.Module):
         return noise
 
     def add_noise(self, waveform:torch.Tensor) -> torch.Tensor:
+        assert self.noise is not None, "Cannot add noise because a noise file was not provided."
         num_repeats = waveform.shape[1] // self.noise.shape[1] + 1
         noise = self.noise.repeat(1,num_repeats)[:, :waveform.shape[1]]
         noise_power = noise.norm(p=2)
@@ -53,7 +55,8 @@ class AudioTrainingPipeline(torch.nn.Module):
         except:
             print("oops")
         waveform = self.preprocess_waveform(waveform)
-        waveform = self.add_noise(waveform)
+        if self.noise is not None:
+            waveform = self.add_noise(waveform)
         spec = self.audio_to_spectrogram(waveform)
 
         # Spectrogram augmentation

train.py

@@ -14,9 +14,8 @@ from models.residual import ResidualDancer, TrainingEnvironment
 import yaml
 from preprocessing.dataset import DanceDataModule, WaveformSongDataset, HuggingFaceWaveformSongDataset
 from torch.utils.data import random_split
-from wakepy import keepawake
 import numpy as np
-from transformers import ASTFeatureExtractor, AutoFeatureExtractor, ASTConfig, AutoModelForAudioClassification
+from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 from argparse import ArgumentParser
 
 
@@ -151,5 +150,4 @@ if __name__ == "__main__":
     config = get_config(args.config)
     training_id = config["global"]["id"]
     train = get_training_fn(training_id)
-    with keepawake():
-        train(config)
+    train(config)