206 lines
6.5 KiB
Python
206 lines
6.5 KiB
Python
import copy
|
|
import numpy as np
|
|
import random
|
|
from collections import defaultdict
|
|
from torch.utils.data.sampler import Sampler, RandomSampler, SequentialSampler
|
|
|
|
|
|
class RandomDomainSampler(Sampler):
|
|
"""Randomly samples N domains each with K images
|
|
to form a minibatch of size N*K.
|
|
|
|
Args:
|
|
data_source (list): list of Datums.
|
|
batch_size (int): batch size.
|
|
n_domain (int): number of domains to sample in a minibatch.
|
|
"""
|
|
|
|
def __init__(self, data_source, batch_size, n_domain):
|
|
self.data_source = data_source
|
|
|
|
# Keep track of image indices for each domain
|
|
self.domain_dict = defaultdict(list)
|
|
for i, item in enumerate(data_source):
|
|
self.domain_dict[item.domain].append(i)
|
|
self.domains = list(self.domain_dict.keys())
|
|
|
|
# Make sure each domain has equal number of images
|
|
if n_domain is None or n_domain <= 0:
|
|
n_domain = len(self.domains)
|
|
assert batch_size % n_domain == 0
|
|
self.n_img_per_domain = batch_size // n_domain
|
|
|
|
self.batch_size = batch_size
|
|
# n_domain denotes number of domains sampled in a minibatch
|
|
self.n_domain = n_domain
|
|
self.length = len(list(self.__iter__()))
|
|
|
|
def __iter__(self):
|
|
domain_dict = copy.deepcopy(self.domain_dict)
|
|
final_idxs = []
|
|
stop_sampling = False
|
|
|
|
while not stop_sampling:
|
|
selected_domains = random.sample(self.domains, self.n_domain)
|
|
|
|
for domain in selected_domains:
|
|
idxs = domain_dict[domain]
|
|
selected_idxs = random.sample(idxs, self.n_img_per_domain)
|
|
final_idxs.extend(selected_idxs)
|
|
|
|
for idx in selected_idxs:
|
|
domain_dict[domain].remove(idx)
|
|
|
|
remaining = len(domain_dict[domain])
|
|
if remaining < self.n_img_per_domain:
|
|
stop_sampling = True
|
|
|
|
return iter(final_idxs)
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
|
|
class SeqDomainSampler(Sampler):
|
|
"""Sequential domain sampler, which randomly samples K
|
|
images from each domain to form a minibatch.
|
|
|
|
Args:
|
|
data_source (list): list of Datums.
|
|
batch_size (int): batch size.
|
|
"""
|
|
|
|
def __init__(self, data_source, batch_size):
|
|
self.data_source = data_source
|
|
|
|
# Keep track of image indices for each domain
|
|
self.domain_dict = defaultdict(list)
|
|
for i, item in enumerate(data_source):
|
|
self.domain_dict[item.domain].append(i)
|
|
self.domains = list(self.domain_dict.keys())
|
|
self.domains.sort()
|
|
|
|
# Make sure each domain has equal number of images
|
|
n_domain = len(self.domains)
|
|
assert batch_size % n_domain == 0
|
|
self.n_img_per_domain = batch_size // n_domain
|
|
|
|
self.batch_size = batch_size
|
|
# n_domain denotes number of domains sampled in a minibatch
|
|
self.n_domain = n_domain
|
|
self.length = len(list(self.__iter__()))
|
|
|
|
def __iter__(self):
|
|
domain_dict = copy.deepcopy(self.domain_dict)
|
|
final_idxs = []
|
|
stop_sampling = False
|
|
|
|
while not stop_sampling:
|
|
for domain in self.domains:
|
|
idxs = domain_dict[domain]
|
|
selected_idxs = random.sample(idxs, self.n_img_per_domain)
|
|
final_idxs.extend(selected_idxs)
|
|
|
|
for idx in selected_idxs:
|
|
domain_dict[domain].remove(idx)
|
|
|
|
remaining = len(domain_dict[domain])
|
|
if remaining < self.n_img_per_domain:
|
|
stop_sampling = True
|
|
|
|
return iter(final_idxs)
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
|
|
class RandomClassSampler(Sampler):
|
|
"""Randomly samples N classes each with K instances to
|
|
form a minibatch of size N*K.
|
|
|
|
Modified from https://github.com/KaiyangZhou/deep-person-reid.
|
|
|
|
Args:
|
|
data_source (list): list of Datums.
|
|
batch_size (int): batch size.
|
|
n_ins (int): number of instances per class to sample in a minibatch.
|
|
"""
|
|
|
|
def __init__(self, data_source, batch_size, n_ins):
|
|
if batch_size < n_ins:
|
|
raise ValueError(
|
|
"batch_size={} must be no less "
|
|
"than n_ins={}".format(batch_size, n_ins)
|
|
)
|
|
|
|
self.data_source = data_source
|
|
self.batch_size = batch_size
|
|
self.n_ins = n_ins
|
|
self.ncls_per_batch = self.batch_size // self.n_ins
|
|
self.index_dic = defaultdict(list)
|
|
for index, item in enumerate(data_source):
|
|
self.index_dic[item.label].append(index)
|
|
self.labels = list(self.index_dic.keys())
|
|
assert len(self.labels) >= self.ncls_per_batch
|
|
|
|
# estimate number of images in an epoch
|
|
self.length = len(list(self.__iter__()))
|
|
|
|
def __iter__(self):
|
|
batch_idxs_dict = defaultdict(list)
|
|
|
|
for label in self.labels:
|
|
idxs = copy.deepcopy(self.index_dic[label])
|
|
if len(idxs) < self.n_ins:
|
|
idxs = np.random.choice(idxs, size=self.n_ins, replace=True)
|
|
random.shuffle(idxs)
|
|
batch_idxs = []
|
|
for idx in idxs:
|
|
batch_idxs.append(idx)
|
|
if len(batch_idxs) == self.n_ins:
|
|
batch_idxs_dict[label].append(batch_idxs)
|
|
batch_idxs = []
|
|
|
|
avai_labels = copy.deepcopy(self.labels)
|
|
final_idxs = []
|
|
|
|
while len(avai_labels) >= self.ncls_per_batch:
|
|
selected_labels = random.sample(avai_labels, self.ncls_per_batch)
|
|
for label in selected_labels:
|
|
batch_idxs = batch_idxs_dict[label].pop(0)
|
|
final_idxs.extend(batch_idxs)
|
|
if len(batch_idxs_dict[label]) == 0:
|
|
avai_labels.remove(label)
|
|
|
|
return iter(final_idxs)
|
|
|
|
def __len__(self):
|
|
return self.length
|
|
|
|
|
|
def build_sampler(
|
|
sampler_type,
|
|
cfg=None,
|
|
data_source=None,
|
|
batch_size=32,
|
|
n_domain=0,
|
|
n_ins=16
|
|
):
|
|
if sampler_type == "RandomSampler":
|
|
return RandomSampler(data_source)
|
|
|
|
elif sampler_type == "SequentialSampler":
|
|
return SequentialSampler(data_source)
|
|
|
|
elif sampler_type == "RandomDomainSampler":
|
|
return RandomDomainSampler(data_source, batch_size, n_domain)
|
|
|
|
elif sampler_type == "SeqDomainSampler":
|
|
return SeqDomainSampler(data_source, batch_size)
|
|
|
|
elif sampler_type == "RandomClassSampler":
|
|
return RandomClassSampler(data_source, batch_size, n_ins)
|
|
|
|
else:
|
|
raise ValueError("Unknown sampler type: {}".format(sampler_type))
|