| 1 | import math |
| 2 | |
| 3 | from torch.utils.data import SequentialSampler |
| 4 | |
| 5 | from f5_tts.model.dataset import DynamicBatchSampler, load_dataset |
| 6 | |
| 7 | |
| 8 | train_dataset = load_dataset("Emilia_ZH_EN", "pinyin") |
| 9 | sampler = SequentialSampler(train_dataset) |
| 10 | |
| 11 | gpus = 8 |
| 12 | batch_size_per_gpu = 38400 |
| 13 | max_samples_per_gpu = 64 |
| 14 | max_updates = 1250000 |
| 15 | |
| 16 | batch_sampler = DynamicBatchSampler( |
| 17 | sampler, |
| 18 | batch_size_per_gpu, |
| 19 | max_samples=max_samples_per_gpu, |
| 20 | random_seed=666, |
| 21 | drop_residual=False, |
| 22 | ) |
| 23 | updates_per_epoch = int(len(batch_sampler) / gpus) |
| 24 | |
| 25 | print( |
| 26 | f"One epoch has {updates_per_epoch} updates if gpus={gpus}, with " |
| 27 | f"batch_size_per_gpu={batch_size_per_gpu} (frames) & " |
| 28 | f"max_samples_per_gpu={max_samples_per_gpu}." |
| 29 | ) |
| 30 | print(f"If gpus={gpus}, for max_updates={max_updates} should set epoch={math.ceil(max_updates / updates_per_epoch)}.") |
| 31 |