Spaces:
Configuration error
Configuration error
File size: 9,466 Bytes
87e21d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 |
# 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])
|