112 lines
4.4 KiB
Python
112 lines
4.4 KiB
Python
from typing import Any, Callable, Dict, List, Optional, Union
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import numpy as np
|
|
import copy
|
|
from torch.utils.data import IterableDataset
|
|
|
|
from longcat_image.utils.dist_utils import get_world_size, get_rank, get_local_rank
|
|
|
|
class MultiResolutionDistributedSampler(torch.utils.data.Sampler):
|
|
def __init__(self,
|
|
batch_size: int,
|
|
dataset: IterableDataset,
|
|
data_resolution_infos: List,
|
|
bucket_info: dict,
|
|
num_replicas: int = None,
|
|
rank: int = None,
|
|
seed: int = 888,
|
|
epoch: int = 0,
|
|
shuffle: bool = True):
|
|
|
|
if not dist.is_available():
|
|
num_replicas = 1
|
|
rank = 0
|
|
else:
|
|
num_replicas = get_world_size()
|
|
rank = get_rank()
|
|
|
|
self.len_items = len(dataset)
|
|
bucket_info = {float(b): bucket_info[b] for b in bucket_info.keys()}
|
|
self.aspect_ratios = np.array(sorted(list(bucket_info.keys())))
|
|
self.resolutions = np.array([bucket_info[aspect] for aspect in self.aspect_ratios])
|
|
|
|
self.batch_size = batch_size
|
|
self.num_replicas = num_replicas
|
|
self.rank = rank
|
|
self.epoch = epoch
|
|
self.shuffle = shuffle
|
|
self.seed = seed
|
|
self.cur_rank_index = []
|
|
self.rng = np.random.RandomState(seed+self.epoch)
|
|
self.global_batch_size = batch_size*num_replicas
|
|
self.data_resolution_infos = np.array(data_resolution_infos, dtype=np.float32)
|
|
print(f'num_replicas {num_replicas}, cur rank {rank}!!!')
|
|
|
|
self.split_to_buckets()
|
|
self.num_samples = len(dataset)//num_replicas
|
|
|
|
def split_to_buckets(self):
|
|
self.buckets = {}
|
|
self._buckets_bak = {}
|
|
data_aspect_ratio = self.data_resolution_infos[:,0]*1.0/self.data_resolution_infos[:, 1]
|
|
bucket_id = np.abs(data_aspect_ratio[:, None] - self.aspect_ratios).argmin(axis=1)
|
|
for i in range(len(self.aspect_ratios)):
|
|
self.buckets[i] = np.where(bucket_id == i)[0]
|
|
self._buckets_bak[i] = np.where(bucket_id == i)[0]
|
|
for k, v in self.buckets.items():
|
|
print(f'bucket {k}, resolutions {self.resolutions[k]}, sampler nums {len(v)}!!!')
|
|
|
|
def get_batch_index(self):
|
|
success_flag = False
|
|
while not success_flag:
|
|
bucket_ids = list(self.buckets.keys())
|
|
bucket_probs = [len(self.buckets[bucket_id]) for bucket_id in bucket_ids]
|
|
bucket_probs = np.array(bucket_probs, dtype=np.float32)
|
|
bucket_probs = bucket_probs / bucket_probs.sum()
|
|
bucket_ids = np.array(bucket_ids, dtype=np.int64)
|
|
chosen_id = int(self.rng.choice(bucket_ids, 1, p=bucket_probs)[0])
|
|
if len(self.buckets[chosen_id]) < self.global_batch_size:
|
|
del self.buckets[chosen_id]
|
|
continue
|
|
batch_data = self.buckets[chosen_id][:self.global_batch_size]
|
|
batch_data = (batch_data, self.resolutions[chosen_id])
|
|
self.buckets[chosen_id] = self.buckets[chosen_id][self.global_batch_size:]
|
|
if len(self.buckets[chosen_id]) == 0:
|
|
del self.buckets[chosen_id]
|
|
success_flag = True
|
|
assert bool(self.buckets), 'There is not enough data in the current epoch.'
|
|
return batch_data
|
|
|
|
def shuffle_bucker_index(self):
|
|
self.rng = np.random.RandomState(self.seed+self.epoch)
|
|
self.buckets = copy.deepcopy(self._buckets_bak)
|
|
for bucket_id in self.buckets.keys():
|
|
self.rng.shuffle(self.buckets[bucket_id])
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
try:
|
|
if len(self.cur_rank_index) == 0:
|
|
global_batch_index, target_resolutions = self.get_batch_index()
|
|
self.cur_rank_index = list(map(
|
|
int, global_batch_index[self.batch_size*self.rank:self.batch_size*(self.rank+1)]))
|
|
self.resolution = list(map(int, target_resolutions))
|
|
data_index = self.cur_rank_index.pop(0)
|
|
return (data_index, self.resolution)
|
|
|
|
except Exception as e:
|
|
self.epoch += 1
|
|
self.shuffle_bucker_index()
|
|
print(f'get error {e}.')
|
|
raise StopIteration
|
|
|
|
def __len__(self):
|
|
return self.num_samples
|
|
|
|
def set_epoch(self, epoch):
|
|
self.epoch = epoch
|