souging/TRELLIS_TextTo3D

1from typing import *2import math3import torch4import numpy as np5from torch.utils.data import Sampler, Dataset, DataLoader, DistributedSampler6import torch.distributed as dist7 8 9def recursive_to_device(10    data: Any,11    device: torch.device,12    non_blocking: bool = False,13) -> Any:14    """15    Recursively move all tensors in a data structure to a device.16    """17    if hasattr(data, "to"):18        return data.to(device, non_blocking=non_blocking)19    elif isinstance(data, (list, tuple)):20        return type(data)(recursive_to_device(d, device, non_blocking) for d in data)21    elif isinstance(data, dict):22        return {k: recursive_to_device(v, device, non_blocking) for k, v in data.items()}23    else:24        return data25 26 27def load_balanced_group_indices(28    load: List[int],29    num_groups: int,30    equal_size: bool = False,31) -> List[List[int]]:32    """33    Split indices into groups with balanced load.34    """35    if equal_size:36        group_size = len(load) // num_groups37    indices = np.argsort(load)[::-1]38    groups = [[] for _ in range(num_groups)]39    group_load = np.zeros(num_groups)40    for idx in indices:41        min_group_idx = np.argmin(group_load)42        groups[min_group_idx].append(idx)43        if equal_size and len(groups[min_group_idx]) == group_size:44            group_load[min_group_idx] = float('inf')45        else:46            group_load[min_group_idx] += load[idx]47    return groups48 49 50def cycle(data_loader: DataLoader) -> Iterator:51    while True:52        for data in data_loader:53            if isinstance(data_loader.sampler, ResumableSampler):54                data_loader.sampler.idx += data_loader.batch_size   # type: ignore[attr-defined]55            yield data56        if isinstance(data_loader.sampler, DistributedSampler):57            data_loader.sampler.epoch += 158        if isinstance(data_loader.sampler, ResumableSampler):59            data_loader.sampler.epoch += 160            data_loader.sampler.idx = 061        62 63class ResumableSampler(Sampler):64    """65    Distributed sampler that is resumable.66 67    Args:68        dataset: Dataset used for sampling.69        rank (int, optional): Rank of the current process within :attr:`num_replicas`.70            By default, :attr:`rank` is retrieved from the current distributed71            group.72        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the73            indices.74        seed (int, optional): random seed used to shuffle the sampler if75            :attr:`shuffle=True`. This number should be identical across all76            processes in the distributed group. Default: ``0``.77        drop_last (bool, optional): if ``True``, then the sampler will drop the78            tail of the data to make it evenly divisible across the number of79            replicas. If ``False``, the sampler will add extra indices to make80            the data evenly divisible across the replicas. Default: ``False``.81    """82 83    def __init__(84        self,85        dataset: Dataset,86        shuffle: bool = True,87        seed: int = 0,88        drop_last: bool = False,89    ) -> None:90        self.dataset = dataset91        self.epoch = 092        self.idx = 093        self.drop_last = drop_last94        self.world_size = dist.get_world_size() if dist.is_initialized() else 195        self.rank = dist.get_rank() if dist.is_initialized() else 096        # If the dataset length is evenly divisible by # of replicas, then there97        # is no need to drop any data, since the dataset will be split equally.98        if self.drop_last and len(self.dataset) % self.world_size != 0:  # type: ignore[arg-type]99            # Split to nearest available length that is evenly divisible.100            # This is to ensure each rank receives the same amount of data when101            # using this Sampler.102            self.num_samples = math.ceil(103                (len(self.dataset) - self.world_size) / self.world_size  # type: ignore[arg-type]104            )105        else:106            self.num_samples = math.ceil(len(self.dataset) / self.world_size)  # type: ignore[arg-type]107        self.total_size = self.num_samples * self.world_size108        self.shuffle = shuffle109        self.seed = seed110 111    def __iter__(self) -> Iterator:112        if self.shuffle:113            # deterministically shuffle based on epoch and seed114            g = torch.Generator()115            g.manual_seed(self.seed + self.epoch)116            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]117        else:118            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]119 120        if not self.drop_last:121            # add extra samples to make it evenly divisible122            padding_size = self.total_size - len(indices)123            if padding_size <= len(indices):124                indices += indices[:padding_size]125            else:126                indices += (indices * math.ceil(padding_size / len(indices)))[127                    :padding_size128                ]129        else:130            # remove tail of data to make it evenly divisible.131            indices = indices[: self.total_size]132        assert len(indices) == self.total_size133 134        # subsample135        indices = indices[self.rank : self.total_size : self.world_size]136        137        # resume from previous state138        indices = indices[self.idx:]139 140        return iter(indices)141 142    def __len__(self) -> int:143        return self.num_samples144 145    def state_dict(self) -> dict[str, int]:146        return {147            'epoch': self.epoch,148            'idx': self.idx,149        }150        151    def load_state_dict(self, state_dict):152        self.epoch = state_dict['epoch']153        self.idx = state_dict['idx']154        155 156class BalancedResumableSampler(ResumableSampler):157    """158    Distributed sampler that is resumable and balances the load among the processes.159 160    Args:161        dataset: Dataset used for sampling.162        rank (int, optional): Rank of the current process within :attr:`num_replicas`.163            By default, :attr:`rank` is retrieved from the current distributed164            group.165        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the166            indices.167        seed (int, optional): random seed used to shuffle the sampler if168            :attr:`shuffle=True`. This number should be identical across all169            processes in the distributed group. Default: ``0``.170        drop_last (bool, optional): if ``True``, then the sampler will drop the171            tail of the data to make it evenly divisible across the number of172            replicas. If ``False``, the sampler will add extra indices to make173            the data evenly divisible across the replicas. Default: ``False``.174    """175 176    def __init__(177        self,178        dataset: Dataset,179        shuffle: bool = True,180        seed: int = 0,181        drop_last: bool = False,182        batch_size: int = 1,183    ) -> None:184        assert hasattr(dataset, 'loads'), 'Dataset must have "loads" attribute to use BalancedResumableSampler'185        super().__init__(dataset, shuffle, seed, drop_last)186        self.batch_size = batch_size187        self.loads = dataset.loads188        189    def __iter__(self) -> Iterator:190        if self.shuffle:191            # deterministically shuffle based on epoch and seed192            g = torch.Generator()193            g.manual_seed(self.seed + self.epoch)194            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]195        else:196            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]197 198        if not self.drop_last:199            # add extra samples to make it evenly divisible200            padding_size = self.total_size - len(indices)201            if padding_size <= len(indices):202                indices += indices[:padding_size]203            else:204                indices += (indices * math.ceil(padding_size / len(indices)))[205                    :padding_size206                ]207        else:208            # remove tail of data to make it evenly divisible.209            indices = indices[: self.total_size]210        assert len(indices) == self.total_size211 212        # balance load among processes213        num_batches = len(indices) // (self.batch_size * self.world_size)214        balanced_indices = []215        for i in range(num_batches):216            start_idx = i * self.batch_size * self.world_size217            end_idx = (i + 1) * self.batch_size * self.world_size218            batch_indices = indices[start_idx:end_idx]219            batch_loads = [self.loads[idx] for idx in batch_indices]220            groups = load_balanced_group_indices(batch_loads, self.world_size, equal_size=True)221            balanced_indices.extend([batch_indices[j] for j in groups[self.rank]])222        223        # resume from previous state224        indices = balanced_indices[self.idx:]225 226        return iter(indices)227