Add ability to merge dataset buckets. Useful when using numerous different datasets.

This commit is contained in:
Jaret Burkett committed 2026-10-08 14:04:30 -06:00
1 parent 0cbac2c1f3
commit a949f21f8e
7 files changed
+270 -95

No files matched your search

+8 -3
View File
@@ -2244,10 +2244,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
ds.cache_tensors_to_disk = True
# load datasets if passed in the root process
if self.datasets is not None:
self.data_loader = get_dataloader_from_datasets(self.datasets, self.train_config.batch_size, self.sd)
self.data_loader = get_dataloader_from_datasets(
self.datasets, self.train_config.batch_size, self.sd,
pool_datasets=self.train_config.pool_datasets,
)
if self.datasets_reg is not None:
self.data_loader_reg = get_dataloader_from_datasets(self.datasets_reg, self.train_config.batch_size,
self.sd)
self.data_loader_reg = get_dataloader_from_datasets(
self.datasets_reg, self.train_config.batch_size, self.sd,
pool_datasets=self.train_config.pool_datasets,
)
flush()
self.last_save_step = self.step_num