From 4d907718f39e2b0f7a0c714c2e3de289e742fc61 Mon Sep 17 00:00:00 2001
From: Carl <415692979@qq.com>
Date: 星期四, 28 三月 2024 13:42:00 +0800
Subject: [PATCH] 修正commit 87b62d68957a2194b017a43b6c2a15424a05a984 引入的英文整句标点预测导致末尾两个单词中间的空格被删除的问题。 (#1556)
---
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