adds parameter to control frequency of checkpoints

src/diffwave/learner.py

@@ -26,7 +26,6 @@ from dataset import from_path, from_gtzan
 from model import DiffWave
 from params import AttrDict
 
-CHECKPOINS_HOP = 10000 # saves checkopints every 10k steps
 
 def _nested_map(struct, map_fn):
   if isinstance(struct, tuple):
@@ -112,7 +111,7 @@ class DiffWaveLearner:
         if self.is_master:
           if self.step % 50 == 0:
             self._write_summary(self.step, features, loss)
-          if self.step % CHECKPOINS_HOP == 0:
+          if self.step % self.params.checkpoints_hop == 0:
             self.save_to_checkpoint()
         self.step += 1
 

src/diffwave/params.py

@@ -53,7 +53,8 @@ params = AttrDict(
     unconditional = True,
     noise_schedule=np.linspace(1e-4, 0.01, 200).tolist(), # last param is num_timesteps
     inference_noise_schedule=[1e-4, 0.001, 0.01, 0.05, 0.2, 0.5], # for fast sampling
-    audio_len = 2000 # length of generated samples
+    audio_len = 2000, # length of generated samples
+    checkpoints_hop = 50000 # how often to save checkpoints
 )