Spaces:
Running on Zero
Running on Zero
| """Dataset sampler from VGGT. | |
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| from __future__ import annotations | |
| import random | |
| import numpy as np | |
| from torch.utils.data import DistributedSampler, RandomSampler, Sampler | |
| class DynamicBatchSampler(Sampler): | |
| """A custom batch sampler. | |
| Dynamically adjusts batch size, aspect ratio, and image number for each | |
| sample. Batches within a sample share the same aspect ratio and image | |
| number. | |
| """ | |
| def __init__( | |
| self, | |
| sampler, | |
| aspect_ratio_range, | |
| image_num_range, | |
| epoch=0, | |
| seed=42, | |
| max_img_per_gpu=48, | |
| ): | |
| """Initializes the dynamic batch sampler. | |
| Args: | |
| sampler: Instance of DynamicDistributedSampler. | |
| aspect_ratio_range: List containing [min_aspect_ratio, max_aspect_ratio]. | |
| image_num_range: List containing [min_images, max_images] per sample. | |
| epoch: Current epoch number. | |
| seed: Random seed for reproducibility. | |
| max_img_per_gpu: Maximum number of images to fit in GPU memory. | |
| """ | |
| self.sampler = sampler | |
| self.aspect_ratio_range = aspect_ratio_range | |
| self.image_num_range = image_num_range | |
| self.rng = random.Random() | |
| # Uniformly sample from the range of possible image numbers | |
| # For any image number, the weight is 1.0 (uniform sampling). You can set any different weights here. | |
| self.image_num_weights = { | |
| num_images: 1.0 | |
| for num_images in range(image_num_range[0], image_num_range[1] + 1) | |
| } | |
| # Possible image numbers, e.g., [2, 3, 4, ..., 24] | |
| self.possible_nums = np.array( | |
| [ | |
| n | |
| for n in self.image_num_weights.keys() | |
| if self.image_num_range[0] <= n <= self.image_num_range[1] | |
| ] | |
| ) | |
| # Normalize weights for sampling | |
| weights = [self.image_num_weights[n] for n in self.possible_nums] | |
| self.normalized_weights = np.array(weights) / sum(weights) | |
| # Maximum image number per GPU | |
| self.max_img_per_gpu = max_img_per_gpu | |
| # Set the epoch for the sampler | |
| self.set_epoch(epoch + seed) | |
| def set_epoch(self, epoch): | |
| """Sets the epoch for this sampler, affecting the random sequence. | |
| Args: | |
| epoch: The epoch number. | |
| """ | |
| self.sampler.set_epoch(epoch) | |
| self.epoch = epoch | |
| self.rng.seed(epoch * 100) | |
| def __iter__(self): | |
| """Yields batches of samples with synchronized dynamic parameters. | |
| Returns: | |
| Iterator yielding batches of indices with associated parameters. | |
| """ | |
| sampler_iterator = iter(self.sampler) | |
| while True: | |
| try: | |
| # Sample random image number and aspect ratio | |
| random_image_num = int( | |
| np.random.choice( | |
| self.possible_nums, p=self.normalized_weights | |
| ) | |
| ) | |
| random_aspect_ratio = round( | |
| self.rng.uniform( | |
| self.aspect_ratio_range[0], self.aspect_ratio_range[1] | |
| ), | |
| 2, | |
| ) | |
| # Update sampler parameters | |
| self.sampler.update_parameters( | |
| aspect_ratio=random_aspect_ratio, | |
| image_num=random_image_num, | |
| ) | |
| # Calculate batch size based on max images per GPU and current | |
| # image number | |
| batch_size = self.max_img_per_gpu / random_image_num | |
| batch_size = np.floor(batch_size).astype(int) | |
| batch_size = max( | |
| 1, batch_size | |
| ) # Ensure batch size is at least 1 | |
| # Collect samples for the current batch | |
| current_batch = [] | |
| for _ in range(batch_size): | |
| try: | |
| item = next( | |
| sampler_iterator | |
| ) # item is (idx, aspect_ratio, image_num) | |
| current_batch.append(item) | |
| except StopIteration: | |
| break # No more samples | |
| if not current_batch: | |
| break # No more data to yield | |
| yield current_batch | |
| except StopIteration: | |
| break # End of sampler's iterator | |
| def __len__(self): | |
| # Return a large dummy length | |
| return 1000000 | |
| class DynamicDistributedSampler(DistributedSampler): | |
| """Extends PyTorch's DistributedSampler. | |
| Include dynamic aspect_ratio and image_num parameters, which can be passed | |
| into the dataset's __getitem__ method. | |
| """ | |
| def __init__( | |
| self, | |
| dataset, | |
| num_replicas: int | None = None, | |
| rank: int | None = None, | |
| shuffle: bool = False, | |
| seed: int = 0, | |
| drop_last: bool = False, | |
| ): | |
| """Init.""" | |
| super().__init__( | |
| dataset, | |
| num_replicas=num_replicas, | |
| rank=rank, | |
| shuffle=shuffle, | |
| seed=seed, | |
| drop_last=drop_last, | |
| ) | |
| self.aspect_ratio = None | |
| self.image_num = None | |
| def __iter__(self): | |
| """Yields a sequence of (index, image_num, aspect_ratio). | |
| Relies on the parent class's logic for shuffling/distributing | |
| the indices across replicas, then attaches extra parameters. | |
| """ | |
| indices_iter = super().__iter__() | |
| for idx in indices_iter: | |
| yield (idx, self.image_num, self.aspect_ratio) | |
| def update_parameters(self, aspect_ratio, image_num): | |
| """Updates dynamic parameters for each new epoch or iteration. | |
| Args: | |
| aspect_ratio: The aspect ratio to set. | |
| image_num: The number of images to set. | |
| """ | |
| self.aspect_ratio = aspect_ratio | |
| self.image_num = image_num | |
| class DynamicSampler(RandomSampler): | |
| """Extends PyTorch's Sampler. | |
| Include dynamic aspect_ratio and image_num parameters, which can be passed | |
| into the dataset's __getitem__ method. | |
| """ | |
| def __init__(self, *args, **kwargs): | |
| """Init.""" | |
| super().__init__(*args, **kwargs) | |
| self.aspect_ratio = None | |
| self.image_num = None | |
| def __iter__(self): | |
| """Yields a sequence of (index, image_num, aspect_ratio). | |
| Relies on the parent class's logic for shuffling/distributing | |
| the indices across replicas, then attaches extra parameters. | |
| """ | |
| indices_iter = super().__iter__() | |
| for idx in indices_iter: | |
| yield (idx, self.image_num, self.aspect_ratio) | |
| def set_epoch(self, epoch): | |
| """Sets the epoch for this sampler, affecting the random sequence. | |
| Args: | |
| epoch: The epoch number. | |
| """ | |
| self.epoch = epoch | |
| def update_parameters(self, aspect_ratio, image_num): | |
| """Updates dynamic parameters for each new epoch or iteration. | |
| Args: | |
| aspect_ratio: The aspect ratio to set. | |
| image_num: The number of images to set. | |
| """ | |
| self.aspect_ratio = aspect_ratio | |
| self.image_num = image_num | |
| def __len__(self) -> int: | |
| return len(self.data_source) | |