make sure soundstream training runs with accelerate

audiolm_pytorch/trainer.py

@@ -39,6 +39,7 @@
 from audiolm_pytorch.utils import AudioConditionerBase
 
 from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
 
 # constants
 
@@ -137,7 +138,9 @@ def __init__(
         force_clear_prev_results = None  # set to True | False to skip the prompt
     ):
         super().__init__()
-        self.accelerator = Accelerator(**accelerate_kwargs)
+
+        kwargs = DistributedDataParallelKwargs(find_unused_parameters = True)
+        self.accelerator = Accelerator(kwargs_handlers = [kwargs], **accelerate_kwargs)
 
         self.soundstream = soundstream
         self.ema_soundstream = EMA(soundstream, beta = ema_beta, update_after_step = ema_update_after_step, update_every = ema_update_every)
@@ -195,14 +198,12 @@ def __init__(
             self.soundstream,
             self.optim,
             self.discr_optim,
-            self.dl,
-            self.valid_dl
+            self.dl
         ) = self.accelerator.prepare(
             self.soundstream,
             self.optim,
             self.discr_optim,
-            self.dl,
-            self.valid_dl
+            self.dl
         )
 
         # prepare the multiscale discriminators with accelerator
@@ -224,7 +225,7 @@ def __init__(
 
         self.results_folder = Path(results_folder)
 
-        if force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
+        if self.is_main and force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
             rmtree(str(self.results_folder))
 
         self.results_folder.mkdir(parents = True, exist_ok = True)
@@ -377,30 +378,37 @@ def train_step(self):
 
         # update exponential moving averaged generator
 
+        self.accelerator.wait_for_everyone()
+
         if self.is_main:
             self.ema_soundstream.update()
 
         # sample results every so often
 
+        self.accelerator.wait_for_everyone()
+
         if self.is_main and not (steps % self.save_results_every):
-            for model, filename in ((self.ema_soundstream.ema_model, f'{steps}.ema'), (self.soundstream, str(steps))):
+            for model, filename in ((self.ema_soundstream.ema_model, f'{steps}.ema'), (self.unwrapped_soundstream, str(steps))):
                 model.eval()
 
                 wave, = next(self.valid_dl_iter)
                 wave = wave.to(device)
 
-                recons = model(wave, return_recons_only = True)
+                with torch.no_grad():
+                    recons = model(wave, return_recons_only = True)
 
                 milestone = steps // self.save_results_every
 
                 for ind, recon in enumerate(recons.unbind(dim = 0)):
                     filename = str(self.results_folder / f'sample_{steps}.flac')
-                    torchaudio.save(filename, recon.cpu().detach(), self.soundstream.target_sample_hz)
+                    torchaudio.save(filename, recon.cpu().detach(), self.unwrapped_soundstream.target_sample_hz)
 
             self.print(f'{steps}: saving to {str(self.results_folder)}')
 
         # save model every so often
 
+        self.accelerator.wait_for_everyone()
+
         if self.is_main and not (steps % self.save_model_every):
             model_path = str(self.results_folder / f'soundstream.{steps}.pt')
             self.save(model_path)
@@ -528,7 +536,7 @@ def __init__(
 
         self.results_folder = Path(results_folder)
 
-        if force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
+        if self.is_main and force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
             rmtree(str(self.results_folder))
 
         self.results_folder.mkdir(parents = True, exist_ok = True)
@@ -763,7 +771,7 @@ def __init__(
 
         self.results_folder = Path(results_folder)
 
-        if force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
+        if self.is_main and force_clear_prev_results is True or (not exists(force_clear_prev_results) and len([*self.results_folder.glob('**/*')]) > 0 and yes_or_no('do you want to clear previous experiment checkpoints and results?')):
             rmtree(str(self.results_folder))
 
         self.results_folder.mkdir(parents = True, exist_ok = True)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'audiolm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.12.4',
+  version = '0.12.5',
   license='MIT',
   description = 'AudioLM - Language Modeling Approach to Audio Generation from Google Research - Pytorch',
   author = 'Phil Wang',