From 93ae7d3fc866038ae9061063bf35b8e13959ba69 Mon Sep 17 00:00:00 2001
From: Can <cansolakoglu130@gmail.com>
Date: Tue, 4 Feb 2025 20:21:59 +0300
Subject: [PATCH 1/2] Enhance DynamicBatchSampler to support epoch-based
 shuffling

---
 src/f5_tts/model/dataset.py | 25 +++++++++++++++++--------
 src/f5_tts/model/trainer.py | 10 +++++++++-
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/f5_tts/model/dataset.py b/src/f5_tts/model/dataset.py
index 10cfd2b..17c9f8b 100644
--- a/src/f5_tts/model/dataset.py
+++ b/src/f5_tts/model/dataset.py
@@ -170,6 +170,7 @@ class DynamicBatchSampler(Sampler[list[int]]):
         in a batch to ensure that the total number of frames are less
         than a certain threshold.
     2.  Make sure the padding efficiency in the batch is high.
+    3.  Shuffle batches each epoch while maintaining reproducibility.
     """
 
     def __init__(
@@ -178,6 +179,8 @@ class DynamicBatchSampler(Sampler[list[int]]):
         self.sampler = sampler
         self.frames_threshold = frames_threshold
         self.max_samples = max_samples
+        self.random_seed = random_seed
+        self.epoch = 0
 
         indices, batches = [], []
         data_source = self.sampler.data_source
@@ -210,17 +213,23 @@ class DynamicBatchSampler(Sampler[list[int]]):
             batches.append(batch)
 
         del indices
-
-        # if want to have different batches between epochs, may just set a seed and log it in ckpt
-        # cuz during multi-gpu training, although the batch on per gpu not change between epochs, the formed general minibatch is different
-        # e.g. for epoch n, use (random_seed + n)
-        random.seed(random_seed)
-        random.shuffle(batches)
-
         self.batches = batches
 
+    def set_epoch(self, epoch: int) -> None:
+        """Sets the epoch for this sampler."""
+        self.epoch = epoch
+
     def __iter__(self):
-        return iter(self.batches)
+        # Use both random_seed and epoch for deterministic but different shuffling per epoch
+        if self.random_seed is not None:
+            g = torch.Generator()
+            g.manual_seed(self.random_seed + self.epoch)
+            # Use PyTorch's random permutation for better reproducibility across PyTorch versions
+            indices = torch.randperm(len(self.batches), generator=g).tolist()
+            batches = [self.batches[i] for i in indices]
+        else:
+            batches = self.batches
+        return iter(batches)
 
     def __len__(self):
         return len(self.batches)
diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py
index 35cd967..b62a767 100644
--- a/src/f5_tts/model/trainer.py
+++ b/src/f5_tts/model/trainer.py
@@ -279,7 +279,11 @@ class Trainer:
             self.accelerator.even_batches = False
             sampler = SequentialSampler(train_dataset)
             batch_sampler = DynamicBatchSampler(
-                sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False
+                sampler, 
+                self.batch_size, 
+                max_samples=self.max_samples, 
+                random_seed=resumable_with_seed,  # This enables reproducible shuffling
+                drop_last=False
             )
             train_dataloader = DataLoader(
                 train_dataset,
@@ -329,6 +333,10 @@ class Trainer:
                 progress_bar_initial = 0
                 current_dataloader = train_dataloader
 
+            # Set epoch for the batch sampler if it exists
+            if hasattr(train_dataloader, 'batch_sampler') and hasattr(train_dataloader.batch_sampler, 'set_epoch'):
+                train_dataloader.batch_sampler.set_epoch(epoch)
+
             progress_bar = tqdm(
                 range(math.ceil(len(train_dataloader) / self.grad_accumulation_steps)),
                 desc=f"Epoch {epoch+1}/{self.epochs}",

From 33e865120c1ea08c39e9fb7b68e2ebd09b37ce58 Mon Sep 17 00:00:00 2001
From: Can <cansolakoglu130@gmail.com>
Date: Tue, 4 Feb 2025 22:20:42 +0300
Subject: [PATCH 2/2] Refactor imports and improve code formatting in dataset
 and trainer modules

---
 src/f5_tts/model/dataset.py |  1 -
 src/f5_tts/model/trainer.py | 10 +++++-----
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/f5_tts/model/dataset.py b/src/f5_tts/model/dataset.py
index 17c9f8b..75eeddd 100644
--- a/src/f5_tts/model/dataset.py
+++ b/src/f5_tts/model/dataset.py
@@ -1,5 +1,4 @@
 import json
-import random
 from importlib.resources import files
 
 import torch
diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py
index b62a767..26970a3 100644
--- a/src/f5_tts/model/trainer.py
+++ b/src/f5_tts/model/trainer.py
@@ -279,11 +279,11 @@ class Trainer:
             self.accelerator.even_batches = False
             sampler = SequentialSampler(train_dataset)
             batch_sampler = DynamicBatchSampler(
-                sampler, 
-                self.batch_size, 
-                max_samples=self.max_samples, 
+                sampler,
+                self.batch_size,
+                max_samples=self.max_samples,
                 random_seed=resumable_with_seed,  # This enables reproducible shuffling
-                drop_last=False
+                drop_last=False,
             )
             train_dataloader = DataLoader(
                 train_dataset,
@@ -334,7 +334,7 @@ class Trainer:
                 current_dataloader = train_dataloader
 
             # Set epoch for the batch sampler if it exists
-            if hasattr(train_dataloader, 'batch_sampler') and hasattr(train_dataloader.batch_sampler, 'set_epoch'):
+            if hasattr(train_dataloader, "batch_sampler") and hasattr(train_dataloader.batch_sampler, "set_epoch"):
                 train_dataloader.batch_sampler.set_epoch(epoch)
 
             progress_bar = tqdm(