diffusion/utils/data_sampler.py

# Copyright (c) OpenMMLab. All rights reserved.
import json
import os
import random
from copy import deepcopy
from random import choice, shuffle
from typing import Sequence

from torch.utils.data import BatchSampler, Dataset, Sampler

from diffusion.utils.logger import get_root_logger


class AspectRatioBatchSampler(BatchSampler):
    """A sampler wrapper for grouping images with similar aspect ratio into a same batch.

    Args:
        sampler (Sampler): Base sampler.
        dataset (Dataset): Dataset providing data information.
        batch_size (int): Size of mini-batch.
        drop_last (bool): If ``True``, the sampler will drop the last batch if
            its size would be less than ``batch_size``.
        aspect_ratios (dict): The predefined aspect ratios.
    """

    def __init__(
        self,
        sampler: Sampler,
        dataset: Dataset,
        batch_size: int,
        aspect_ratios: dict,
        drop_last: bool = False,
        config=None,
        valid_num=0,  # take as valid aspect-ratio when sample number >= valid_num
        hq_only=False,
        cache_file=None,
        caching=False,
        **kwargs,
    ) -> None:
        if not isinstance(sampler, Sampler):
            raise TypeError(f"sampler should be an instance of ``Sampler``, but got {sampler}")
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise ValueError(f"batch_size should be a positive integer value, but got batch_size={batch_size}")

        self.sampler = sampler
        self.dataset = dataset
        self.batch_size = batch_size
        self.aspect_ratios = aspect_ratios
        self.drop_last = drop_last
        self.hq_only = hq_only
        self.config = config
        self.caching = caching
        self.cache_file = cache_file
        self.order_check_pass = False

        self.ratio_nums_gt = kwargs.get("ratio_nums", None)
        assert self.ratio_nums_gt, "ratio_nums_gt must be provided."
        self._aspect_ratio_buckets = {ratio: [] for ratio in aspect_ratios.keys()}
        self.current_available_bucket_keys = [str(k) for k, v in self.ratio_nums_gt.items() if v >= valid_num]

        logger = (
            get_root_logger() if config is None else get_root_logger(os.path.join(config.work_dir, "train_log.log"))
        )
        logger.warning(
            f"Using valid_num={valid_num} in config file. Available {len(self.current_available_bucket_keys)} aspect_ratios: {self.current_available_bucket_keys}"
        )

        self.data_all = {} if caching else None
        if os.path.exists(cache_file):
            logger.info(f"Loading cached file for multi-scale training: {cache_file}")
            try:
                self.cached_idx = json.load(open(cache_file))
            except:
                logger.info(f"Failed loading: {cache_file}")
                self.cached_idx = {}
        else:
            logger.info(f"No cached file is found, dataloader is slow: {cache_file}")
            self.cached_idx = {}
        self.exist_ids = len(self.cached_idx)

    def __iter__(self) -> Sequence[int]:
        for idx in self.sampler:
            data_info, closest_ratio = self._get_data_info_and_ratio(idx)
            if not data_info:
                continue

            bucket = self._aspect_ratio_buckets[closest_ratio]
            bucket.append(idx)
            # yield a batch of indices in the same aspect ratio group
            if len(bucket) == self.batch_size:
                self._update_cache(bucket)
                yield bucket[:]
                del bucket[:]

        for bucket in self._aspect_ratio_buckets.values():
            while bucket:
                if not self.drop_last or len(bucket) == self.batch_size:
                    yield bucket[:]
                del bucket[:]

    def _get_data_info_and_ratio(self, idx):
        str_idx = str(idx)
        if self.caching:
            if str_idx in self.cached_idx:
                return self.cached_idx[str_idx], self.cached_idx[str_idx]["closest_ratio"]
            data_info = self.dataset.get_data_info(int(idx))
            if data_info is None or (
                self.hq_only and "version" in data_info and data_info["version"] not in ["high_quality"]
            ):
                return None, None
            closest_ratio = self._get_closest_ratio(data_info["height"], data_info["width"])
            self.data_all[str_idx] = {
                "height": data_info["height"],
                "width": data_info["width"],
                "closest_ratio": closest_ratio,
                "key": data_info["key"],
            }
            return data_info, closest_ratio
        else:
            if self.cached_idx:
                if self.cached_idx.get(str_idx):
                    if not self.order_check_pass or random.random() < 0.01:
                        # Ensure the cached dataset is in the same order as the original tar file
                        self._order_check(str_idx)
                    closest_ratio = self.cached_idx[str_idx]["closest_ratio"]
                    return self.cached_idx[str_idx], closest_ratio

            data_info = self.dataset.get_data_info(int(idx))
            if data_info is None or (
                self.hq_only and "version" in data_info and data_info["version"] not in ["high_quality"]
            ):
                return None, None
            closest_ratio = self._get_closest_ratio(data_info["height"], data_info["width"])

            return data_info, closest_ratio

    def _get_closest_ratio(self, height, width):
        ratio = height / width
        return min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio))

    def _order_check(self, str_idx):
        ori_data = self.cached_idx[str_idx]
        real_key = self.dataset.get_data_info(int(str_idx))["key"]
        assert real_key and ori_data["key"] == real_key, ValueError(
            f"index: {str_idx}, real key: {real_key} ori key: {ori_data['key']}"
        )
        self.order_check_pass = True

    def _update_cache(self, bucket):
        if self.caching:
            for idx in bucket:
                if str(idx) in self.cached_idx:
                    continue
                self.cached_idx[str(idx)] = self.data_all.pop(str(idx))


class BalancedAspectRatioBatchSampler(AspectRatioBatchSampler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Assign samples to each bucket
        self.ratio_nums_gt = kwargs.get("ratio_nums", None)
        assert self.ratio_nums_gt
        self._aspect_ratio_buckets = {float(ratio): [] for ratio in self.aspect_ratios.keys()}
        self.original_buckets = {}
        self.current_available_bucket_keys = [k for k, v in self.ratio_nums_gt.items() if v >= 3000]
        self.all_available_keys = deepcopy(self.current_available_bucket_keys)
        self.exhausted_bucket_keys = []
        self.total_batches = len(self.sampler) // self.batch_size
        self._aspect_ratio_count = {}
        for k in self.all_available_keys:
            self._aspect_ratio_count[float(k)] = 0
            self.original_buckets[float(k)] = []
        logger = get_root_logger(os.path.join(self.config.work_dir, "train_log.log"))
        logger.warning(
            f"Available {len(self.current_available_bucket_keys)} aspect_ratios: {self.current_available_bucket_keys}"
        )

    def __iter__(self) -> Sequence[int]:
        i = 0
        for idx in self.sampler:
            data_info = self.dataset.get_data_info(idx)
            height, width = data_info["height"], data_info["width"]
            ratio = height / width
            closest_ratio = float(min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio)))
            if closest_ratio not in self.all_available_keys:
                continue
            if self._aspect_ratio_count[closest_ratio] < self.ratio_nums_gt[closest_ratio]:
                self._aspect_ratio_count[closest_ratio] += 1
                self._aspect_ratio_buckets[closest_ratio].append(idx)
                self.original_buckets[closest_ratio].append(idx)  # Save the original samples for each bucket
            if not self.current_available_bucket_keys:
                self.current_available_bucket_keys, self.exhausted_bucket_keys = self.exhausted_bucket_keys, []

            if closest_ratio not in self.current_available_bucket_keys:
                continue
            key = closest_ratio
            bucket = self._aspect_ratio_buckets[key]
            if len(bucket) == self.batch_size:
                yield bucket[: self.batch_size]
                del bucket[: self.batch_size]
                i += 1
                self.exhausted_bucket_keys.append(key)
                self.current_available_bucket_keys.remove(key)

        for _ in range(self.total_batches - i):
            key = choice(self.all_available_keys)
            bucket = self._aspect_ratio_buckets[key]
            if len(bucket) >= self.batch_size:
                yield bucket[: self.batch_size]
                del bucket[: self.batch_size]

                # If a bucket is exhausted
                if not bucket:
                    self._aspect_ratio_buckets[key] = deepcopy(self.original_buckets[key][:])
                    shuffle(self._aspect_ratio_buckets[key])
            else:
                self._aspect_ratio_buckets[key] = deepcopy(self.original_buckets[key][:])
                shuffle(self._aspect_ratio_buckets[key])