mirror of
https://github.com/hpcaitech/Open-Sora.git
synced 2026-05-12 04:59:35 +02:00
119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
import random
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.distributed import ProcessGroup
|
|
from torch.distributed.distributed_c10d import _get_default_group
|
|
from torch.utils.data import DataLoader
|
|
|
|
from .sampler import StatefulDistributedSampler, VariableVideoBatchSampler
|
|
|
|
|
|
def prepare_dataloader(
|
|
dataset,
|
|
batch_size,
|
|
shuffle=False,
|
|
seed=1024,
|
|
drop_last=False,
|
|
pin_memory=False,
|
|
num_workers=0,
|
|
process_group: Optional[ProcessGroup] = None,
|
|
distributed=True,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Prepare a dataloader for distributed training. The dataloader will be wrapped by
|
|
`torch.utils.data.DataLoader` and `StatefulDistributedSampler`.
|
|
|
|
|
|
Args:
|
|
dataset (`torch.utils.data.Dataset`): The dataset to be loaded.
|
|
shuffle (bool, optional): Whether to shuffle the dataset. Defaults to False.
|
|
seed (int, optional): Random worker seed for sampling, defaults to 1024.
|
|
add_sampler: Whether to add ``DistributedDataParallelSampler`` to the dataset. Defaults to True.
|
|
drop_last (bool, optional): Set to True to drop the last incomplete batch, if the dataset size
|
|
is not divisible by the batch size. If False and the size of dataset is not divisible by
|
|
the batch size, then the last batch will be smaller, defaults to False.
|
|
pin_memory (bool, optional): Whether to pin memory address in CPU memory. Defaults to False.
|
|
num_workers (int, optional): Number of worker threads for this dataloader. Defaults to 0.
|
|
kwargs (dict): optional parameters for ``torch.utils.data.DataLoader``, more details could be found in
|
|
`DataLoader <https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html#DataLoader>`_.
|
|
|
|
Returns:
|
|
:class:`torch.utils.data.DataLoader`: A DataLoader used for training or testing.
|
|
"""
|
|
_kwargs = kwargs.copy()
|
|
if distributed:
|
|
process_group = process_group or _get_default_group()
|
|
sampler = StatefulDistributedSampler(
|
|
dataset,
|
|
num_replicas=process_group.size(),
|
|
rank=process_group.rank(),
|
|
shuffle=shuffle,
|
|
)
|
|
else:
|
|
sampler = None
|
|
|
|
# Deterministic dataloader
|
|
def seed_worker(worker_id):
|
|
worker_seed = seed
|
|
np.random.seed(worker_seed)
|
|
torch.manual_seed(worker_seed)
|
|
random.seed(worker_seed)
|
|
|
|
return DataLoader(
|
|
dataset,
|
|
batch_size=batch_size,
|
|
sampler=sampler,
|
|
worker_init_fn=seed_worker,
|
|
drop_last=drop_last,
|
|
pin_memory=pin_memory,
|
|
num_workers=num_workers,
|
|
**_kwargs,
|
|
)
|
|
|
|
|
|
def prepare_variable_dataloader(
|
|
dataset,
|
|
batch_size,
|
|
bucket_config,
|
|
shuffle=False,
|
|
seed=1024,
|
|
drop_last=False,
|
|
pin_memory=False,
|
|
num_workers=0,
|
|
process_group=None,
|
|
num_bucket_build_workers=1,
|
|
**kwargs,
|
|
):
|
|
_kwargs = kwargs.copy()
|
|
process_group = process_group or _get_default_group()
|
|
batch_sampler = VariableVideoBatchSampler(
|
|
dataset,
|
|
bucket_config,
|
|
num_replicas=process_group.size(),
|
|
rank=process_group.rank(),
|
|
shuffle=shuffle,
|
|
seed=seed,
|
|
drop_last=drop_last,
|
|
verbose=True,
|
|
num_bucket_build_workers=num_bucket_build_workers,
|
|
)
|
|
|
|
# Deterministic dataloader
|
|
def seed_worker(worker_id):
|
|
worker_seed = seed
|
|
np.random.seed(worker_seed)
|
|
torch.manual_seed(worker_seed)
|
|
random.seed(worker_seed)
|
|
|
|
return torch.utils.data.DataLoader(
|
|
dataset,
|
|
batch_sampler=batch_sampler,
|
|
worker_init_fn=seed_worker,
|
|
pin_memory=pin_memory,
|
|
num_workers=num_workers,
|
|
**_kwargs,
|
|
)
|