From d1b1fdd520cd33f0fc297eaa6ee6f451a85781cd Mon Sep 17 00:00:00 2001
From: zhifu gao <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 28 三月 2024 10:17:50 +0800
Subject: [PATCH] Dev gzf new (#1555)
---
funasr/datasets/audio_datasets/samplers.py | 15 +++++++++++----
1 files changed, 11 insertions(+), 4 deletions(-)
diff --git a/funasr/datasets/audio_datasets/samplers.py b/funasr/datasets/audio_datasets/samplers.py
index 01f5e6a..b4fb846 100644
--- a/funasr/datasets/audio_datasets/samplers.py
+++ b/funasr/datasets/audio_datasets/samplers.py
@@ -2,6 +2,7 @@
import numpy as np
import logging
import math
+import random
import torch.distributed as dist
from torch.utils.data import DistributedSampler
from torch.utils.data import BatchSampler, Sampler
@@ -322,17 +323,21 @@
self.shuffle = shuffle and is_training
self.drop_last = drop_last
- self.total_size = len(self.dataset)
- # self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
+ # self.total_size = len(self.dataset)
+ self.num_samples = int(math.ceil(self.total_size / self.num_replicas))
self.epoch = 0
self.sort_size = sort_size * num_replicas
self.max_token_length = kwargs.get("max_token_length", 2048)
self.length_scale_source = kwargs.get("length_scale_source", 1.0)
+ super().__init__(dataset, num_replicas=num_replicas, rank=rank,
+ shuffle=shuffle, drop_last=drop_last)
def __iter__(self):
if self.shuffle:
g = torch.Generator()
g.manual_seed(self.epoch)
+ random.seed(self.epoch)
+
indices = torch.randperm(len(self.dataset), generator=g).tolist()
else:
indices = list(range(len(self.dataset)))
@@ -362,8 +367,10 @@
# Ensure each rank gets the same number of batches, duplicate data if needed
batches_per_rank = math.ceil(len(buffer_batches) / self.num_replicas)
total_batches_needed = batches_per_rank * self.num_replicas
- buffer_batches.extend(buffer_batches[:total_batches_needed - len(buffer_batches)])
-
+
+ extra_batches = total_batches_needed - len(buffer_batches)
+ buffer_batches += random.choices(buffer_batches, k=extra_batches)
+
# Evenly distribute batches from buffer_batches to each rank
rank_batches = [[] for _ in range(self.num_replicas)]
for i, batch in enumerate(buffer_batches):
--
Gitblit v1.9.1